<commit_before>#
# Copyright (c) 2010 rPath, Inc.
#
# All Rights Reserved
#

import logging
from StringIO import StringIO

from smartform import descriptor
from smartform import descriptor_errors

from conary import conaryclient, versions
from conary import trove as conarytrove
from conary.errors import RepositoryError

from django.core.exceptions import ObjectDoesNotExist

from mint.django_rest import timeutils
from mint.django_rest.rbuilder.inventory import models
from mint.django_rest.rbuilder.manager import basemanager
from mint.django_rest.rbuilder.projects.models import Project, ProjectVersion
from mint.rest.errors import ProductNotFound, ProductVersionNotFound

log = logging.getLogger(__name__)
exposed = basemanager.exposed

class VersionManager(basemanager.BaseManager):
    """
    Class encapsulating all logic around versions, available updates, etc.
    """

    def __init__(self, *args, **kwargs):
        basemanager.BaseManager.__init__(self, *args, **kwargs)
        self._cclient = None

    def get_software_versions(self, system):
        pass

    @classmethod
    def _flavor(cls, flavor):
        if flavor is None:
            return ''
        return str(flavor)

    def getStages(self, hostname, majorVersionName):
        stages = self.restDb.getProductVersionStages(hostname,
            majorVersionName)
        return stages.stages

    def setStage(self, system, trove):
        if not trove.is_top_level:
            return

        hostname = trove.getHost()

        label = trove.version.conaryVersion.trailingLabel()
        tag = label.getLabel()

        # NOTE: This assumes that repoName and the tag always match.
        repoName = hostname.split('.')[0]
        rparts = repoName.split('-')
        tparts = tag.split('-')
        majorVersionName = tparts[len(rparts)]

        stages = self.getStages(hostname,
            majorVersionName)
        stage = [s for s in stages \
            if s.label == trove.version.label]
        if not stage:
            return

        stage = stage[0]
        try:
            project = Project.objects.get(repository_hostname=hostname)
            majorVersion = ProjectVersion.objects.get(project=project,
                name=majorVersionName)
        except ObjectDoesNotExist:
            return

        stage, created = models.Stage.objects.get_or_create(project_branch=majorVersion,
            label=trove.version.label, name=stage.name, project=project)

        system.project_branch_stage = stage
        system.project_branch = majorVersion
        system.project = project

    def trove_from_nvf(self, nvf):
        n, v, f = conaryclient.cmdline.parseTroveSpec(nvf)
        f = self._flavor(f)
        thawed_v = versions.ThawVersion(v)
        version = models.Version()
        version.fromConaryVersion(thawed_v)
        version.flavor = f

        trove = models.Trove()
        trove.name = n
        trove.version = version
        trove.flavor = f
        return trove

    @classmethod
    def _trove(cls, trove):
        """
        If the trove is new, save it to the db, otherwise return the existing
        one
        """
        # First, make sure the flavor is part of the version object
        if not trove.version.flavor:
            trove.version.flavor = trove.flavor
        created, version = models.Version.objects.load_or_create(trove.version)
        trove.version = version
        created, trove = models.Trove.objects.load_or_create(trove)
        # Need to call save again so that is_top_level gets reset
        trove.save()
        return trove

    def cache_available_update(self, nvf, update_nvf):
        trove = self.trove_from_nvf(nvf)
        update_trove = self.trove_from_nvf(update_nvf)
        available_update, created = models.AvailableUpdate.get_or_create(
            trove=trove, trove_available_update=update_trove)
        available_update.save()

    def get_conary_client(self):
        if self._cclient is None:
            self._cclient = self.restDb.productMgr.reposMgr.getUserClient()
        return self._cclient

    def _checkCacheExpired(self, trove):
        one_day = timeutils.timedelta(1)
        return (trove.last_available_update_refresh + one_day) < timeutils.now()

    # FIXME: is this still used?  -- MPD
    def set_available_updates(self, trove, force=False):

        # Hack to make sure utc is set as the timezone on
        # last_available_update_refresh.
        if trove.last_available_update_refresh is not None:
            trove.last_available_update_refresh = \
                trove.last_available_update_refresh.replace(tzinfo=timeutils.TZUTC)

        if force or \
           trove.last_available_update_refresh is None or \
           self._checkCacheExpired(trove):

            refreshed = self.refresh_available_updates(trove)
            if refreshed:
                trove.last_available_update_refresh = timeutils.now()
                trove.save()

    # NOTE: this system should no longer be used, but not positive -- MPD
    def refresh_available_updates(self, trove):
        self.cclient = self.get_conary_client()
        # trvName and trvVersion are str's, trvFlavor is a
        # conary.deps.deps.Flavor.
        trvName = trove.name
        trvVersion = trove.version.conaryVersion
        trvFlavor = trove.getFlavor()
        trvLabel = trove.getLabel()

        # Search the label for the trove of the top level item.  It should
        # only (hopefully) return 1 result.
        try:
            troves = self.cclient.repos.findTroves(trvLabel,
                [(trvName, trvVersion, trvFlavor)])
        except RepositoryError, e:
            log.error("Error contacting repository to look for available " + \
                "updates for %s=%s[%s]" % (trvName, trvLabel, trvFlavor))
            log.error(e)
            return False
        except ProductNotFound, e:
            log.error("Permission error querying repository for %s=%s[%s]" \
                % (trvName, trvLabel, trvFlavor))
            log.error(e)
            return False
        assert(len(troves) == 1)

        trove.available_updates.clear()
        trove.out_of_date = False

        # findTroves returns a {} with keys of (name, version, flavor), values
        # of [(name, repoVersion, repoFlavor)], where repoVersion and
        # repoFlavor are rich objects with the repository metadata.
        repoVersion = troves[(trvName, trvVersion, trvFlavor)][0][1]
        repoFlavors = [f[0][2] for f in troves.values()]
        # We only asked for 1 flavor, only 1 should be returned.
        assert(len(repoFlavors) == 1)

        # getTroveVersionList searches a repository (NOT by label), for a
        # given name/flavor combination.
        allVersions = self.cclient.repos.getTroveVersionList(
            trvVersion.getHost(), {trvName:repoFlavors})
        # We only asked for 1 name/flavor, so we should have only gotten 1
        # back.
        assert(len(allVersions) == 1)
        # getTroveVersionList returns a dict with keys of name, values of
        # {version: [flavors]}.
        allVersions = allVersions[trvName]

        availableVersions = {}
        for v, fs in allVersions.iteritems():
            # getTroveVersionList doesn't search by label, so we need to
            # compare the results to the label we're interested in.
            if v.trailingLabel() == trvLabel:

                # Check that at least one of the flavors found satisfies the
                # flavor we're interested in.
                satisfiedFlavors = []
                for f in fs:
                    # XXX: do we want to use flavor or repoFlavor here?
                    # XXX: do we want to use stronglySatisfies here?
                    if f.satisfies(trvFlavor):
                        satisfiedFlavors.append(f)
                if satisfiedFlavors:
                    availableVersions[v] = satisfiedFlavors

        if availableVersions:
            for ver, fs in availableVersions.iteritems():
                for flv in fs:
                    new_version = models.Version()
                    new_version.fromConaryVersion(ver)
                    new_version.flavor = str(flv)
                    created, new_version = \
                        models.Version.objects.load_or_create(new_version)
                    new_version.save()
                    trove.available_updates.add(new_version)

                    # Set out of date flag on True if there are newer versions
                    # available.
                    if ver > repoVersion:
                        trove.out_of_date = True

        # Always add the current version as an available update, this is so
        # that remediation will work.
        trove.available_updates.add(trove.version)

        return True

    # FIMXE: this from old mint/rest/api
    def _getTroveConfigDescriptor(self, name, version, flavor):

        repos = self.get_conary_client().repos
        trvList = repos.getTroves([(name, version, flavor)])

        referencedByDefault = []
        for trv in trvList:
            referencedByDefault += [ nvf for nvf, byDefault, strongRef in
                trv.iterTroveListInfo() if byDefault ]

        # Get properties sorted by package name.
        properties = repos.getTroveInfo(conarytrove._TROVEINFO_TAG_PROPERTIES,
            sorted(referencedByDefault, cmp=lambda x, y: cmp(x[0], y[0])))

        configFields = []
        for propSet in properties:
            if propSet is None:
                continue
            for property in propSet.iter():
                xml = property.definition()
                desc = descriptor.BaseDescriptor()

                try:
                    desc.parseStream(StringIO(xml))

                # Ignore any descriptors that don't parse.
                except descriptor_errors.Error:
                    continue

                configFields.extend(desc.getDataFields())

        return configFields

    # this code originally from mint/rest/api
    def _getConfigDescriptor(self, name, version, flavor):

        # FIXME: import descriptor
        desc = descriptor.ConfigurationDescriptor()
        desc.setDisplayName('Configuration Descriptor')
        desc.addDescription('Configuration Descriptor')

        newFields = self._getTroveConfigDescriptor(name, version, flavor)
        if not newFields:
            return ''

        fields = desc.getDataFields()
        fields.extend(newFields)

        out = StringIO()
        desc.serialize(out, validate=False)
        out.seek(0)

        return out.read()

    @exposed
    def getConfigurationDescriptor(self, system):
        """
        Generate config descriptor for all top level items on a system.
        """

        if system.latest_survey is None:
            return '<configuration></configuration>'

        # find the appliance group from the survey
        # what if multiple top levels with config descriptors?
        # UI doesn't support, so not worrying about it for now
        packages = system.latest_survey.conary_packages
        for conary_package in packages.all():
            info = conary_package.conary_package_info
            name = info.name
            if name.startswith("group-") and name.find("-appliance") != -1:
                #try:
                thawed_v = versions.ThawVersion(info.version)
                version = models.Version()
                version.fromConaryVersion(thawed_v)
                version.flavor = info.flavor

                #except versions.ParseError, e:
                #    raise errors.InvalidVersion("Error parsing version %s: %s" %
                #        (version, str(e)))
                res = self._getConfigDescriptor(info.name, version, info.flavor)
                return res

        # shouldn't ever get here unless you migrated to something weird, in which
        # case (FIXME) just present the empty one and maybe log?
        raise Exception("could not find group-X-appliance")

<commit_msg>cleanup version/flavor parsing<commit_after>#
# Copyright (c) 2010 rPath, Inc.
#
# All Rights Reserved
#

import logging
from StringIO import StringIO

from smartform import descriptor
from smartform import descriptor_errors

from conary.deps import deps
from conary import conaryclient, versions
from conary import trove as conarytrove
from conary.errors import RepositoryError

from django.core.exceptions import ObjectDoesNotExist

from mint.django_rest import timeutils
from mint.django_rest.rbuilder.inventory import models
from mint.django_rest.rbuilder.manager import basemanager
from mint.django_rest.rbuilder.projects.models import Project, ProjectVersion
from mint.rest.errors import ProductNotFound, ProductVersionNotFound

log = logging.getLogger(__name__)
exposed = basemanager.exposed

class VersionManager(basemanager.BaseManager):
    """
    Class encapsulating all logic around versions, available updates, etc.
    """

    def __init__(self, *args, **kwargs):
        basemanager.BaseManager.__init__(self, *args, **kwargs)
        self._cclient = None

    def get_software_versions(self, system):
        pass

    @classmethod
    def _flavor(cls, flavor):
        if flavor is None:
            return ''
        return str(flavor)

    def getStages(self, hostname, majorVersionName):
        stages = self.restDb.getProductVersionStages(hostname,
            majorVersionName)
        return stages.stages

    def setStage(self, system, trove):
        if not trove.is_top_level:
            return

        hostname = trove.getHost()

        label = trove.version.conaryVersion.trailingLabel()
        tag = label.getLabel()

        # NOTE: This assumes that repoName and the tag always match.
        repoName = hostname.split('.')[0]
        rparts = repoName.split('-')
        tparts = tag.split('-')
        majorVersionName = tparts[len(rparts)]

        stages = self.getStages(hostname,
            majorVersionName)
        stage = [s for s in stages \
            if s.label == trove.version.label]
        if not stage:
            return

        stage = stage[0]
        try:
            project = Project.objects.get(repository_hostname=hostname)
            majorVersion = ProjectVersion.objects.get(project=project,
                name=majorVersionName)
        except ObjectDoesNotExist:
            return

        stage, created = models.Stage.objects.get_or_create(project_branch=majorVersion,
            label=trove.version.label, name=stage.name, project=project)

        system.project_branch_stage = stage
        system.project_branch = majorVersion
        system.project = project

    def trove_from_nvf(self, nvf):
        n, v, f = conaryclient.cmdline.parseTroveSpec(nvf)
        f = self._flavor(f)
        thawed_v = versions.ThawVersion(v)
        version = models.Version()
        version.fromConaryVersion(thawed_v)
        version.flavor = f

        trove = models.Trove()
        trove.name = n
        trove.version = version
        trove.flavor = f
        return trove

    @classmethod
    def _trove(cls, trove):
        """
        If the trove is new, save it to the db, otherwise return the existing
        one
        """
        # First, make sure the flavor is part of the version object
        if not trove.version.flavor:
            trove.version.flavor = trove.flavor
        created, version = models.Version.objects.load_or_create(trove.version)
        trove.version = version
        created, trove = models.Trove.objects.load_or_create(trove)
        # Need to call save again so that is_top_level gets reset
        trove.save()
        return trove

    def cache_available_update(self, nvf, update_nvf):
        trove = self.trove_from_nvf(nvf)
        update_trove = self.trove_from_nvf(update_nvf)
        available_update, created = models.AvailableUpdate.get_or_create(
            trove=trove, trove_available_update=update_trove)
        available_update.save()

    def get_conary_client(self):
        if self._cclient is None:
            self._cclient = self.restDb.productMgr.reposMgr.getUserClient()
        return self._cclient

    def _checkCacheExpired(self, trove):
        one_day = timeutils.timedelta(1)
        return (trove.last_available_update_refresh + one_day) < timeutils.now()

    # FIXME: is this still used?  -- MPD
    def set_available_updates(self, trove, force=False):

        # Hack to make sure utc is set as the timezone on
        # last_available_update_refresh.
        if trove.last_available_update_refresh is not None:
            trove.last_available_update_refresh = \
                trove.last_available_update_refresh.replace(tzinfo=timeutils.TZUTC)

        if force or \
           trove.last_available_update_refresh is None or \
           self._checkCacheExpired(trove):

            refreshed = self.refresh_available_updates(trove)
            if refreshed:
                trove.last_available_update_refresh = timeutils.now()
                trove.save()

    # NOTE: this system should no longer be used, but not positive -- MPD
    def refresh_available_updates(self, trove):
        self.cclient = self.get_conary_client()
        # trvName and trvVersion are str's, trvFlavor is a
        # conary.deps.deps.Flavor.
        trvName = trove.name
        trvVersion = trove.version.conaryVersion
        trvFlavor = trove.getFlavor()
        trvLabel = trove.getLabel()

        # Search the label for the trove of the top level item.  It should
        # only (hopefully) return 1 result.
        try:
            troves = self.cclient.repos.findTroves(trvLabel,
                [(trvName, trvVersion, trvFlavor)])
        except RepositoryError, e:
            log.error("Error contacting repository to look for available " + \
                "updates for %s=%s[%s]" % (trvName, trvLabel, trvFlavor))
            log.error(e)
            return False
        except ProductNotFound, e:
            log.error("Permission error querying repository for %s=%s[%s]" \
                % (trvName, trvLabel, trvFlavor))
            log.error(e)
            return False
        assert(len(troves) == 1)

        trove.available_updates.clear()
        trove.out_of_date = False

        # findTroves returns a {} with keys of (name, version, flavor), values
        # of [(name, repoVersion, repoFlavor)], where repoVersion and
        # repoFlavor are rich objects with the repository metadata.
        repoVersion = troves[(trvName, trvVersion, trvFlavor)][0][1]
        repoFlavors = [f[0][2] for f in troves.values()]
        # We only asked for 1 flavor, only 1 should be returned.
        assert(len(repoFlavors) == 1)

        # getTroveVersionList searches a repository (NOT by label), for a
        # given name/flavor combination.
        allVersions = self.cclient.repos.getTroveVersionList(
            trvVersion.getHost(), {trvName:repoFlavors})
        # We only asked for 1 name/flavor, so we should have only gotten 1
        # back.
        assert(len(allVersions) == 1)
        # getTroveVersionList returns a dict with keys of name, values of
        # {version: [flavors]}.
        allVersions = allVersions[trvName]

        availableVersions = {}
        for v, fs in allVersions.iteritems():
            # getTroveVersionList doesn't search by label, so we need to
            # compare the results to the label we're interested in.
            if v.trailingLabel() == trvLabel:

                # Check that at least one of the flavors found satisfies the
                # flavor we're interested in.
                satisfiedFlavors = []
                for f in fs:
                    # XXX: do we want to use flavor or repoFlavor here?
                    # XXX: do we want to use stronglySatisfies here?
                    if f.satisfies(trvFlavor):
                        satisfiedFlavors.append(f)
                if satisfiedFlavors:
                    availableVersions[v] = satisfiedFlavors

        if availableVersions:
            for ver, fs in availableVersions.iteritems():
                for flv in fs:
                    new_version = models.Version()
                    new_version.fromConaryVersion(ver)
                    new_version.flavor = str(flv)
                    created, new_version = \
                        models.Version.objects.load_or_create(new_version)
                    new_version.save()
                    trove.available_updates.add(new_version)

                    # Set out of date flag on True if there are newer versions
                    # available.
                    if ver > repoVersion:
                        trove.out_of_date = True

        # Always add the current version as an available update, this is so
        # that remediation will work.
        trove.available_updates.add(trove.version)

        return True

    def _getTroveConfigDescriptor(self, name, version, flavor):
        repos = self.get_conary_client().repos
        trvList = repos.getTroves([(name, version, flavor)])

        referencedByDefault = []
        for trv in trvList:
            referencedByDefault += [ nvf for nvf, byDefault, strongRef in
                trv.iterTroveListInfo() if byDefault ]

        # Get properties sorted by package name.
        properties = repos.getTroveInfo(conarytrove._TROVEINFO_TAG_PROPERTIES,
            sorted(referencedByDefault, cmp=lambda x, y: cmp(x[0], y[0])))

        configFields = []
        for propSet in properties:
            if propSet is None:
                continue
            for property in propSet.iter():
                xml = property.definition()
                desc = descriptor.BaseDescriptor()

                try:
                    desc.parseStream(StringIO(xml))

                # Ignore any descriptors that don't parse.
                except descriptor_errors.Error:
                    continue

                configFields.extend(desc.getDataFields())

        return configFields

    def _getConfigDescriptor(self, name, version, flavor):
        desc = descriptor.ConfigurationDescriptor()
        desc.setDisplayName('Configuration Descriptor')
        desc.addDescription('Configuration Descriptor')

        newFields = self._getTroveConfigDescriptor(name, version, flavor)
        if not newFields:
            return ''

        fields = desc.getDataFields()
        fields.extend(newFields)

        out = StringIO()
        desc.serialize(out, validate=False)
        out.seek(0)

        return out.read()

    @exposed
    def getConfigurationDescriptor(self, system):
        """
        Generate config descriptor for all top level items on a system.
        """

        if system.latest_survey is None:
            return '<configuration></configuration>'

        # find the appliance group from the survey
        # what if multiple top levels with config descriptors?
        # UI doesn't support, so not worrying about it for now
        packages = system.latest_survey.conary_packages
        for conary_package in packages.all():
            info = conary_package.conary_package_info
            name = info.name
            version = versions.ThawVersion(info.version)
            flavor = deps.parseFlavor(info.flavor)
            if name.startswith("group-") and name.find("-appliance") != -1:
                return self._getConfigDescriptor(info.name, version, flavor)

        # shouldn't ever get here unless you migrated to something weird, in
        # which case (FIXME) just present the empty one and maybe log?
        raise Exception("could not find group-X-appliance")
<|endoftext|>


<commit_before>#
# Copyright (c) 2010 rPath, Inc.
#
# All Rights Reserved
#

import logging
from StringIO import StringIO

from smartform import descriptor
from smartform import descriptor_errors

from conary import conaryclient, versions
from conary import trove as conarytrove
from conary.errors import RepositoryError

from django.core.exceptions import ObjectDoesNotExist

from mint.django_rest import timeutils
from mint.django_rest.rbuilder.inventory import models
from mint.django_rest.rbuilder.manager import basemanager
from mint.django_rest.rbuilder.projects.models import Project, ProjectVersion
from mint.rest.errors import ProductNotFound, ProductVersionNotFound

log = logging.getLogger(__name__)
exposed = basemanager.exposed

class VersionManager(basemanager.BaseManager):
    """
    Class encapsulating all logic around versions, available updates, etc.
    """

    def __init__(self, *args, **kwargs):
        basemanager.BaseManager.__init__(self, *args, **kwargs)
        self._cclient = None

    def get_software_versions(self, system):
        pass

    def delete_installed_software(self, system):
        system.installed_software.all().delete()

    def _diffVersions(self, system, new_versions):
        oldInstalled = dict((x.getNVF(), x)
            for x in system.installed_software.all())

        # Do the delta
        toAdd = []
        newInstalled = []
        for new_version in new_versions:
            if isinstance(new_version, basestring):
                trove = self.trove_from_nvf(new_version)
            else:
                trove = new_version
            nvf = trove.getNVF()
            isInst = oldInstalled.pop(nvf, None)
            if isInst is None:
                toAdd.append(self._trove(trove))
            newInstalled.append(nvf)

        return oldInstalled, newInstalled, toAdd

    @exposed
    def setInstalledSoftware(self, system, new_versions):
        oldInstalled, newInstalled, toAdd = \
            self._diffVersions(system, new_versions)
        for trove in oldInstalled.itervalues():
            system.installed_software.remove(trove)
        for trove in toAdd:
            system.installed_software.add(trove)
            try:
                self.setStage(system, trove)
            except ProductVersionNotFound:
                system.project = None
        for trove in system.installed_software.all():
            self.set_available_updates(trove, force=True)
        system.save()

    @classmethod
    def _flavor(cls, flavor):
        if flavor is None:
            return ''
        return str(flavor)

    def getStages(self, hostname, majorVersionName):
        stages = self.restDb.getProductVersionStages(hostname,
            majorVersionName)
        return stages.stages

    def setStage(self, system, trove):
        if not trove.is_top_level:
            return

        hostname = trove.getHost()

        label = trove.version.conaryVersion.trailingLabel()
        tag = label.getLabel()

        # NOTE: This assumes that repoName and the tag always match.
        repoName = hostname.split('.')[0]
        rparts = repoName.split('-')
        tparts = tag.split('-')
        majorVersionName = tparts[len(rparts)]

        stages = self.getStages(hostname,
            majorVersionName)
        stage = [s for s in stages \
            if s.label == trove.version.label]
        if not stage:
            return

        stage = stage[0]
        try:
            project = Project.objects.get(repository_hostname=hostname)
            majorVersion = ProjectVersion.objects.get(project=project,
                name=majorVersionName)
        except ObjectDoesNotExist:
            return

        stage, created = models.Stage.objects.get_or_create(project_branch=majorVersion,
            label=trove.version.label, name=stage.name, project=project)

        system.project_branch_stage = stage
        system.project_branch = majorVersion
        system.project = project

    @exposed
    def updateInstalledSoftware(self, system_id, new_versions):
        system = models.System.objects.get(pk=system_id)
        troveSpecs = ["%s=%s[%s]" % x.getNVF()
            for x in new_versions ]
        if troveSpecs:
            msg = "Initiating software update to: %s" % (
                ', '.join(troveSpecs), )
        else:
            msg = "Initiating software update, deleting everything"
        self.mgr.log_system(system, msg)
        oldInstalled, newInstalled, toAdd = \
            self._diffVersions(system, new_versions)
        sources = []
        for nvf in newInstalled:
            n, v, f = nvf
            sources.append("%s=%s[%s]" % (n, str(v), self._flavor(f)))
        self.mgr.scheduleSystemApplyUpdateEvent(system, sources)

    def trove_from_nvf(self, nvf):
        n, v, f = conaryclient.cmdline.parseTroveSpec(nvf)
        f = self._flavor(f)

        thawed_v = versions.ThawVersion(v)
        version = models.Version()
        version.fromConaryVersion(thawed_v)
        version.flavor = f

        trove = models.Trove()
        trove.name = n
        trove.version = version
        trove.flavor = f
        return trove

    @classmethod
    def _trove(cls, trove):
        """
        If the trove is new, save it to the db, otherwise return the existing
        one
        """
        # First, make sure the flavor is part of the version object
        if not trove.version.flavor:
            trove.version.flavor = trove.flavor
        created, version = models.Version.objects.load_or_create(trove.version)
        trove.version = version
        created, trove = models.Trove.objects.load_or_create(trove)
        # Need to call save again so that is_top_level gets reset
        trove.save()
        return trove

    def cache_available_update(self, nvf, update_nvf):
        trove = self.trove_from_nvf(nvf)
        update_trove = self.trove_from_nvf(update_nvf)
        available_update, created = models.AvailableUpdate.get_or_create(
            trove=trove, trove_available_update=update_trove)
        available_update.save()

    @exposed
    def refreshCachedUpdates(self, name, label):
        troves = models.Trove.objects.filter(name=name,
            version__label=label)
        for trove in troves:
            self.set_available_updates(trove, force=True)

    def get_conary_client(self):
        if self._cclient is None:
            self._cclient = self.restDb.productMgr.reposMgr.getUserClient()
        return self._cclient

    def _checkCacheExpired(self, trove):
        one_day = timeutils.timedelta(1)
        return (trove.last_available_update_refresh + one_day) < timeutils.now()

    @exposed
    def set_available_updates(self, trove, force=False):

        # Hack to make sure utc is set as the timezone on
        # last_available_update_refresh.
        if trove.last_available_update_refresh is not None:
            trove.last_available_update_refresh = \
                trove.last_available_update_refresh.replace(tzinfo=timeutils.TZUTC)

        if force or \
           trove.last_available_update_refresh is None or \
           self._checkCacheExpired(trove):

            refreshed = self.refresh_available_updates(trove)
            if refreshed:
                trove.last_available_update_refresh = timeutils.now()
                trove.save()

    def refresh_available_updates(self, trove):
        self.cclient = self.get_conary_client()
        # trvName and trvVersion are str's, trvFlavor is a
        # conary.deps.deps.Flavor.
        trvName = trove.name
        trvVersion = trove.version.conaryVersion
        trvFlavor = trove.getFlavor()
        trvLabel = trove.getLabel()

        # Search the label for the trove of the top level item.  It should
        # only (hopefully) return 1 result.
        try:
            troves = self.cclient.repos.findTroves(trvLabel,
                [(trvName, trvVersion, trvFlavor)])
        except RepositoryError, e:
            log.error("Error contacting repository to look for available " + \
                "updates for %s=%s[%s]" % (trvName, trvLabel, trvFlavor))
            log.error(e)
            return False
        except ProductNotFound, e:
            log.error("Permission error querying repository for %s=%s[%s]" \
                % (trvName, trvLabel, trvFlavor))
            log.error(e)
            return False
        assert(len(troves) == 1)

        trove.available_updates.clear()
        trove.out_of_date = False

        # findTroves returns a {} with keys of (name, version, flavor), values
        # of [(name, repoVersion, repoFlavor)], where repoVersion and
        # repoFlavor are rich objects with the repository metadata.
        repoVersion = troves[(trvName, trvVersion, trvFlavor)][0][1]
        repoFlavors = [f[0][2] for f in troves.values()]
        # We only asked for 1 flavor, only 1 should be returned.
        assert(len(repoFlavors) == 1)

        # getTroveVersionList searches a repository (NOT by label), for a
        # given name/flavor combination.
        allVersions = self.cclient.repos.getTroveVersionList(
            trvVersion.getHost(), {trvName:repoFlavors})
        # We only asked for 1 name/flavor, so we should have only gotten 1
        # back.
        assert(len(allVersions) == 1)
        # getTroveVersionList returns a dict with keys of name, values of
        # {version: [flavors]}.
        allVersions = allVersions[trvName]

        availableVersions = {}
        for v, fs in allVersions.iteritems():
            # getTroveVersionList doesn't search by label, so we need to
            # compare the results to the label we're interested in.
            if v.trailingLabel() == trvLabel:

                # Check that at least one of the flavors found satisfies the
                # flavor we're interested in.
                satisfiedFlavors = []
                for f in fs:
                    # XXX: do we want to use flavor or repoFlavor here?
                    # XXX: do we want to use stronglySatisfies here?
                    if f.satisfies(trvFlavor):
                        satisfiedFlavors.append(f)
                if satisfiedFlavors:
                    availableVersions[v] = satisfiedFlavors

        if availableVersions:
            for ver, fs in availableVersions.iteritems():
                for flv in fs:
                    new_version = models.Version()
                    new_version.fromConaryVersion(ver)
                    new_version.flavor = str(flv)
                    created, new_version = \
                        models.Version.objects.load_or_create(new_version)
                    new_version.save()
                    trove.available_updates.add(new_version)

                    # Set out of date flag on True if there are newer versions
                    # available.
                    if ver > repoVersion:
                        trove.out_of_date = True

        # Always add the current version as an available update, this is so
        # that remediation will work.
        trove.available_updates.add(trove.version)

        return True

    @exposed
    def getConfigurationDescriptor(self, system):
        """
        Generate config descriptor for all top level items on a system.
        """

        desc = descriptor.SystemConfigurationDescriptor()
        desc.setDisplayName('Configuration Descriptor')
        desc.addDescription('Configuration Descriptor')

        fields = desc.getDataFields()
        for trove in system.installed_software.all():
            fields.extend(self._getTroveConfigDescriptor(trove))

        out = StringIO()
        desc.serialize(out, validate=False)
        out.seek(0)

        return out.read()

    def _getTroveConfigDescriptor(self, trove):
        client = self.get_conary_client()
        repos = client.getRepos()
        n, v, f = trove.getNVF()

        trvList = repos.getTroves([(n, v, f)])

        referencedByDefault = []
        for trv in trvList:
            referencedByDefault += [ nvf for nvf, byDefault, strongRef in
                trv.iterTroveListInfo() if byDefault ]

        # Get properties sorted by package name.
        properties = repos.getTroveInfo(conarytrove._TROVEINFO_TAG_PROPERTIES,
            sorted(referencedByDefault, cmp=lambda x, y: cmp(x[0], y[0])))

        configFields = []
        for propSet in properties:
            if propSet is None:
                continue
            for property in propSet.iter():
                xml = property.definition()
                desc = descriptor.BaseDescriptor()

                try:
                    desc.parseStream(StringIO(xml))

                # Ignore any descriptors that don't parse.
                except descriptor_errors.Error:
                    continue

                configFields.extend(desc.getDataFields())

        return configFields
<commit_msg>Branch merge from p5<commit_after>#
# Copyright (c) 2010 rPath, Inc.
#
# All Rights Reserved
#

import logging
from StringIO import StringIO

from smartform import descriptor
from smartform import descriptor_errors

from conary import conaryclient, versions
from conary import trove as conarytrove
from conary.errors import RepositoryError

from django.core.exceptions import ObjectDoesNotExist

from mint.django_rest import timeutils
from mint.django_rest.rbuilder.inventory import models
from mint.django_rest.rbuilder.manager import basemanager
from mint.django_rest.rbuilder.projects.models import Project, ProjectVersion
from mint.rest.errors import ProductNotFound, ProductVersionNotFound

log = logging.getLogger(__name__)
exposed = basemanager.exposed

class VersionManager(basemanager.BaseManager):
    """
    Class encapsulating all logic around versions, available updates, etc.
    """

    def __init__(self, *args, **kwargs):
        basemanager.BaseManager.__init__(self, *args, **kwargs)
        self._cclient = None

    def get_software_versions(self, system):
        pass

    def delete_installed_software(self, system):
        system.installed_software.all().delete()

    def _diffVersions(self, system, new_versions):
        oldInstalled = dict((x.getNVF(), x)
            for x in system.installed_software.all())

        # Do the delta
        toAdd = []
        newInstalled = []
        for new_version in new_versions:
            if isinstance(new_version, basestring):
                trove = self.trove_from_nvf(new_version)
            else:
                trove = new_version
            nvf = trove.getNVF()
            isInst = oldInstalled.pop(nvf, None)
            if isInst is None:
                toAdd.append(self._trove(trove))
            newInstalled.append(nvf)

        return oldInstalled, newInstalled, toAdd

    @exposed
    def setInstalledSoftware(self, system, new_versions):
        oldInstalled, newInstalled, toAdd = \
            self._diffVersions(system, new_versions)
        for trove in oldInstalled.itervalues():
            system.installed_software.remove(trove)
        for trove in toAdd:
            system.installed_software.add(trove)
            try:
                self.setStage(system, trove)
            except ProductVersionNotFound:
                system.project = None
        for trove in system.installed_software.all():
            self.set_available_updates(trove, force=True)
        system.save()

    @classmethod
    def _flavor(cls, flavor):
        if flavor is None:
            return ''
        return str(flavor)

    def getStages(self, hostname, majorVersionName):
        stages = self.restDb.getProductVersionStages(hostname,
            majorVersionName)
        return stages.stages

    def setStage(self, system, trove):
        if not trove.is_top_level:
            return

        hostname = trove.getHost()

        label = trove.version.conaryVersion.trailingLabel()
        tag = label.getLabel()

        # NOTE: This assumes that repoName and the tag always match.
        repoName = hostname.split('.')[0]
        rparts = repoName.split('-')
        tparts = tag.split('-')
        majorVersionName = tparts[len(rparts)]

        stages = self.getStages(hostname,
            majorVersionName)
        stage = [s for s in stages \
            if s.label == trove.version.label]
        if not stage:
            return

        stage = stage[0]
        try:
            project = Project.objects.get(repository_hostname=hostname)
            majorVersion = ProjectVersion.objects.get(project=project,
                name=majorVersionName)
        except ObjectDoesNotExist:
            return

        stage, created = models.Stage.objects.get_or_create(project_branch=majorVersion,
            label=trove.version.label, name=stage.name, project=project)

        system.project_branch_stage = stage
        system.project_branch = majorVersion
        system.project = project

    @exposed
    def updateInstalledSoftware(self, system_id, new_versions):
        system = models.System.objects.get(pk=system_id)
        troveSpecs = ["%s=%s[%s]" % x.getNVF()
            for x in new_versions ]
        if troveSpecs:
            msg = "Initiating software update to: %s" % (
                ', '.join(troveSpecs), )
        else:
            msg = "Initiating software update, deleting everything"
        self.mgr.log_system(system, msg)
        oldInstalled, newInstalled, toAdd = \
            self._diffVersions(system, new_versions)
        sources = []
        for nvf in newInstalled:
            n, v, f = nvf
            sources.append("%s=%s[%s]" % (n, str(v), self._flavor(f)))
        self.mgr.scheduleSystemApplyUpdateEvent(system, sources)

    def trove_from_nvf(self, nvf):
        n, v, f = conaryclient.cmdline.parseTroveSpec(nvf)
        f = self._flavor(f)

        thawed_v = versions.ThawVersion(v)
        version = models.Version()
        version.fromConaryVersion(thawed_v)
        version.flavor = f

        trove = models.Trove()
        trove.name = n
        trove.version = version
        trove.flavor = f
        return trove

    @classmethod
    def _trove(cls, trove):
        """
        If the trove is new, save it to the db, otherwise return the existing
        one
        """
        # First, make sure the flavor is part of the version object
        if not trove.version.flavor:
            trove.version.flavor = trove.flavor
        created, version = models.Version.objects.load_or_create(trove.version)
        trove.version = version
        created, trove = models.Trove.objects.load_or_create(trove)
        # Need to call save again so that is_top_level gets reset
        trove.save()
        return trove

    def cache_available_update(self, nvf, update_nvf):
        trove = self.trove_from_nvf(nvf)
        update_trove = self.trove_from_nvf(update_nvf)
        available_update, created = models.AvailableUpdate.get_or_create(
            trove=trove, trove_available_update=update_trove)
        available_update.save()

    @exposed
    def refreshCachedUpdates(self, name, label):
        troves = models.Trove.objects.filter(name=name,
            version__label=label)
        for trove in troves:
            self.set_available_updates(trove, force=True)

    def get_conary_client(self):
        if self._cclient is None:
            self._cclient = self.restDb.productMgr.reposMgr.getUserClient()
        return self._cclient

    def _checkCacheExpired(self, trove):
        one_day = timeutils.timedelta(1)
        return (trove.last_available_update_refresh + one_day) < timeutils.now()

    @exposed
    def set_available_updates(self, trove, force=False):

        # Hack to make sure utc is set as the timezone on
        # last_available_update_refresh.
        if trove.last_available_update_refresh is not None:
            trove.last_available_update_refresh = \
                trove.last_available_update_refresh.replace(tzinfo=timeutils.TZUTC)

        if force or \
           trove.last_available_update_refresh is None or \
           self._checkCacheExpired(trove):

            refreshed = self.refresh_available_updates(trove)
            if refreshed:
                trove.last_available_update_refresh = timeutils.now()
                trove.save()

    def refresh_available_updates(self, trove):
        self.cclient = self.get_conary_client()
        # trvName and trvVersion are str's, trvFlavor is a
        # conary.deps.deps.Flavor.
        trvName = trove.name
        trvVersion = trove.version.conaryVersion
        trvFlavor = trove.getFlavor()
        trvLabel = trove.getLabel()

        # Search the label for the trove of the top level item.  It should
        # only (hopefully) return 1 result.
        try:
            troves = self.cclient.repos.findTroves(trvLabel,
                [(trvName, trvVersion, trvFlavor)])
        except RepositoryError, e:
            log.error("Error contacting repository to look for available " + \
                "updates for %s=%s[%s]" % (trvName, trvLabel, trvFlavor))
            log.error(e)
            return False
        except ProductNotFound, e:
            log.error("Permission error querying repository for %s=%s[%s]" \
                % (trvName, trvLabel, trvFlavor))
            log.error(e)
            return False
        assert(len(troves) == 1)

        trove.available_updates.clear()
        trove.out_of_date = False

        # findTroves returns a {} with keys of (name, version, flavor), values
        # of [(name, repoVersion, repoFlavor)], where repoVersion and
        # repoFlavor are rich objects with the repository metadata.
        repoVersion = troves[(trvName, trvVersion, trvFlavor)][0][1]
        repoFlavors = [f[0][2] for f in troves.values()]
        # We only asked for 1 flavor, only 1 should be returned.
        assert(len(repoFlavors) == 1)

        # getTroveVersionList searches a repository (NOT by label), for a
        # given name/flavor combination.
        allVersions = self.cclient.repos.getTroveVersionList(
            trvVersion.getHost(), {trvName:repoFlavors})
        # We only asked for 1 name/flavor, so we should have only gotten 1
        # back.
        assert(len(allVersions) == 1)
        # getTroveVersionList returns a dict with keys of name, values of
        # {version: [flavors]}.
        allVersions = allVersions[trvName]

        availableVersions = {}
        for v, fs in allVersions.iteritems():
            # getTroveVersionList doesn't search by label, so we need to
            # compare the results to the label we're interested in.
            if v.trailingLabel() == trvLabel:

                # Check that at least one of the flavors found satisfies the
                # flavor we're interested in.
                satisfiedFlavors = []
                for f in fs:
                    # XXX: do we want to use flavor or repoFlavor here?
                    # XXX: do we want to use stronglySatisfies here?
                    if f.satisfies(trvFlavor):
                        satisfiedFlavors.append(f)
                if satisfiedFlavors:
                    availableVersions[v] = satisfiedFlavors

        if availableVersions:
            for ver, fs in availableVersions.iteritems():
                for flv in fs:
                    new_version = models.Version()
                    new_version.fromConaryVersion(ver)
                    new_version.flavor = str(flv)
                    created, new_version = \
                        models.Version.objects.load_or_create(new_version)
                    new_version.save()
                    trove.available_updates.add(new_version)

                    # Set out of date flag on True if there are newer versions
                    # available.
                    if ver > repoVersion:
                        trove.out_of_date = True

        # Always add the current version as an available update, this is so
        # that remediation will work.
        trove.available_updates.add(trove.version)

        return True

    @exposed
    def getConfigurationDescriptor(self, system):
        """
        Generate config descriptor for all top level items on a system.
        """

        desc = descriptor.SystemConfigurationDescriptor()

        fields = desc.getDataFields()
        for trove in system.installed_software.all():
            fields.extend(self._getTroveConfigDescriptor(trove))

        out = StringIO()
        desc.serialize(out, validate=False)
        out.seek(0)

        return out.read()

    def _getTroveConfigDescriptor(self, trove):
        client = self.get_conary_client()
        repos = client.getRepos()
        n, v, f = trove.getNVF()

        trvList = repos.getTroves([(n, v, f)])

        referencedByDefault = []
        for trv in trvList:
            referencedByDefault += [ nvf for nvf, byDefault, strongRef in
                trv.iterTroveListInfo() if byDefault ]

        # Get properties sorted by package name.
        properties = repos.getTroveInfo(conarytrove._TROVEINFO_TAG_PROPERTIES,
            sorted(referencedByDefault, cmp=lambda x, y: cmp(x[0], y[0])))

        configFields = []
        for propSet in properties:
            if propSet is None:
                continue
            for property in propSet.iter():
                xml = property.definition()
                desc = descriptor.BaseDescriptor()

                try:
                    desc.parseStream(StringIO(xml))

                # Ignore any descriptors that don't parse.
                except descriptor_errors.Error:
                    continue

                configFields.extend(desc.getDataFields())

        return configFields
<|endoftext|>


<commit_before>import sys
import os
import settings
import difflib
import utils
import shutil
import json
import filecmp

"""
This script compares the generated CSV coverage files with the ones in the codebase.
"""

artifacts_worflow_name = "Check framework coverage changes"

def check_file_exists(file):
    if not os.path.exists(file):
        print("Expected file '" + file + "' doesn't exist.", file=sys.stderr)
        return False
    return True

def ignore_line_ending(ch):
    return difflib.IS_CHARACTER_JUNK(ch, ws=" \r\n")

def compare_files(file1, file2):
    messages = compare_files_str(file1, file2)
    if messages == "":
        return True

    print(messages, end="", file=sys.stderr)

    return False

def compare_files_str(file1, file2):
    diff = difflib.ndiff(open(file1).readlines(),
                         open(file2).readlines(), None, ignore_line_ending)
    ret = ""
    for line in diff:
        if line.startswith("+") or line.startswith("-"):
            ret += line

    return ret

def write_diff_for_run(output_file, repo, run_id):
    folder1 = "out_base"
    folder2 = "out_merge"
    try:
        utils.subprocess_run(["gh", "run", "download", "--repo", repo, "--name",
                              "csv-framework-coverage-base", "--dir", folder1, str(run_id)])
        utils.subprocess_run(["gh", "run", "download", "--repo", repo, "--name",
                              "csv-framework-coverage-merge", "--dir", folder2, str(run_id)])
        utils.subprocess_run(["gh", "run", "download", "--repo", repo, "--name",
                              "pr", "--dir", "pr", str(run_id)])

        compare_folders(folder1, folder2, output_file)
    finally:
        if os.path.isdir(folder1):
            shutil.rmtree(folder1)

        if os.path.isdir(folder2):
            shutil.rmtree(folder2)

def get_comment_text(output_file, repo, run_id):
    size = os.path.getsize(output_file)
    if size == 0:
        print("No difference in the coverage reports")
        return

    comment = ":warning: The head of this PR and the base branch were compared for differences in the framework coverage reports. " + \
        "The generated reports are available in the [artifacts of this workflow run](https://github.com/" + repo + "/actions/runs/" + str(run_id) + "). " + \
        "The differences will be picked up by the nightly job after the PR gets merged. "

    if size < 2000:
        print("There's a small change in the CSV framework coverage reports")
        comment += "The following differences were found: \n\n"
        with open(output_file, 'r') as file:
            comment += file.read()
    else:
        print("There's a large change in the CSV framework coverage reports")
        comment += "The differences can be found in the " + \
            output_file + " artifact of this job."

    return comment

def comment_pr(output_file, repo, run_id):
    """
    Generates coverage diff produced by the changes in the current PR. If the diff is not empty, then post it as a comment.
    If a workflow run produces the same diff as the directly preceeding one, then don't post a comment.
    """

    # Store diff for current run
    write_diff_for_run(output_file, repo, run_id)

    try:
        with open("pr/NR") as file:
            pr_number = int(file.read())
    finally:
        if os.path.isdir("pr"):
            shutil.rmtree("pr")

    # Try storing diff for previous run:
    prev_output_file = "prev_" + output_file
    try:
        prev_run_id = get_previous_run_id(repo, run_id, pr_number)
        write_diff_for_run(prev_output_file, repo, prev_run_id)

        if filecmp.cmp(output_file, prev_output_file, shallow=False):
            print("Previous run " + str(prev_run_id) +
                  " resulted in the same diff, so not commenting again.")
            return
        else:
            print("Diff of previous run " +
                  str(prev_run_id) + " differs, commenting.")
    except Exception:
        # this is not mecessarily a failure, it can also mean that there was no previous run yet.
        print("Couldn't generate diff for previous run:", sys.exc_info()[1])
    finally:
        if os.path.isfile(prev_output_file):
            os.remove(prev_output_file)

    comment = get_comment_text(output_file, repo, run_id)
    post_comment(comment, repo, pr_number)

def post_comment(comment, repo, pr_number):
    print("Posting comment to PR #" + str(pr_number))
    utils.subprocess_run(["gh", "pr", "comment", str(pr_number),
                         "--repo", repo, "--body", comment])

def compare_folders(folder1, folder2, output_file):
    """
    Compares the contents of two folders and writes the differences to the output file.
    """

    languages = ['java']

    return_md = ""

    for lang in languages:
        expected_files = ""

        generated_output_rst = settings.generated_output_rst.format(
            language=lang)
        generated_output_csv = settings.generated_output_csv.format(
            language=lang)

        # check if files exist in both folder1 and folder 2
        if not check_file_exists(folder1 + "/" + generated_output_rst):
            expected_files += "- " + generated_output_rst + \
                " doesn't exist in folder " + folder1 + "\n"
        if not check_file_exists(folder2 + "/" + generated_output_rst):
            expected_files += "- " + generated_output_rst + \
                " doesn't exist in folder " + folder2 + "\n"
        if not check_file_exists(folder1 + "/" + generated_output_csv):
            expected_files += "- " + generated_output_csv + \
                " doesn't exist in folder " + folder1 + "\n"
        if not check_file_exists(folder2 + "/" + generated_output_csv):
            expected_files += "- " + generated_output_csv + \
                " doesn't exist in folder " + folder2 + "\n"

        if expected_files != "":
            print("Expected files are missing", file=sys.stderr)
            return_md += "\n### " + lang + "\n\n#### Expected files are missing for " + \
                lang + "\n" + expected_files + "\n"
            continue

        # compare contents of files
        cmp1 = compare_files_str(
            folder1 + "/" + generated_output_rst, folder2 + "/" + generated_output_rst)
        cmp2 = compare_files_str(
            folder1 + "/" + generated_output_csv, folder2 + "/" + generated_output_csv)

        if cmp1 != "" or cmp2 != "":
            print("Generated file contents are not matching", file=sys.stderr)
            return_md += "\n### " + lang + "\n\n#### Generated file changes for " + \
                lang + "\n\n"
            if cmp1 != "":
                return_md += "- Changes to " + generated_output_rst + \
                    ":\n```diff\n" + cmp1 + "```\n\n"
            if cmp2 != "":
                return_md += "- Changes to " + generated_output_csv + \
                    ":\n```diff\n" + cmp2 + "```\n\n"

    with open(output_file, 'w', newline='') as out:
        out.write(return_md)

def get_previous_run_id(repo, run_id, pr_number):
    """
    Gets the previous run id for a given workflow run, considering that the previous workflow run needs to come from the same PR.
    """

    # Get branch and repo from run:
    this_run = utils.subprocess_check_output(["gh", "api", "-X", "GET", "repos/" + repo + "/actions/runs/" + str(
        run_id), "--jq", "{ head_branch: .head_branch, head_repository: .head_repository.full_name }"])

    this_run = json.loads(this_run)
    pr_branch = this_run["head_branch"]
    pr_repo = this_run["head_repository"]

    # Get all previous runs that match branch, repo and workflow name:
    ids = utils.subprocess_check_output(["gh", "api", "-X", "GET", "repos/" + repo + "/actions/runs", "-f", "event=pull_request", "-f", "status=success", "-f", "name=\"" + artifacts_worflow_name + "\"", "--jq",
                                        "[.workflow_runs.[] | select(.head_branch==\"" + pr_branch + "\" and .head_repository.full_name==\"" + pr_repo + "\") | { created_at: .created_at, run_id: .id}] | sort_by(.created_at) | reverse | [.[].run_id]"])

    ids = json.loads(ids)
    if ids[0] != run_id:
        raise Exception("Expected to find " + str(run_id) +
                        " in the list of matching runs.")

    for previous_run_id in ids[1:]:
        utils.subprocess_run(["gh", "run", "download", "--repo", repo,
                             "--name", "pr", "--dir", "prev_run_pr", str(previous_run_id)])

        try:
            with open("prev_run_pr/NR") as file:
                prev_pr_number = int(file.read())
                print("PR number: " + str(prev_pr_number))
        finally:
            if os.path.isdir("prev_run_pr"):
                shutil.rmtree("prev_run_pr")

        # the previous run needs to be coming from the same PR:
        if pr_number == prev_pr_number:
            return previous_run_id

    raise Exception("Couldn't find previous run.")

output_file = sys.argv[1]
repo = sys.argv[2]
run_id = sys.argv[3]

comment_pr(output_file, repo, run_id)
<commit_msg>Fix string vs int ID comparison<commit_after>import sys
import os
import settings
import difflib
import utils
import shutil
import json
import filecmp

"""
This script compares the generated CSV coverage files with the ones in the codebase.
"""

artifacts_worflow_name = "Check framework coverage changes"

def check_file_exists(file):
    if not os.path.exists(file):
        print("Expected file '" + file + "' doesn't exist.", file=sys.stderr)
        return False
    return True

def ignore_line_ending(ch):
    return difflib.IS_CHARACTER_JUNK(ch, ws=" \r\n")

def compare_files(file1, file2):
    messages = compare_files_str(file1, file2)
    if messages == "":
        return True

    print(messages, end="", file=sys.stderr)

    return False

def compare_files_str(file1, file2):
    diff = difflib.ndiff(open(file1).readlines(),
                         open(file2).readlines(), None, ignore_line_ending)
    ret = ""
    for line in diff:
        if line.startswith("+") or line.startswith("-"):
            ret += line

    return ret

def write_diff_for_run(output_file, repo, run_id):
    folder1 = "out_base"
    folder2 = "out_merge"
    try:
        utils.subprocess_run(["gh", "run", "download", "--repo", repo, "--name",
                              "csv-framework-coverage-base", "--dir", folder1, str(run_id)])
        utils.subprocess_run(["gh", "run", "download", "--repo", repo, "--name",
                              "csv-framework-coverage-merge", "--dir", folder2, str(run_id)])
        utils.subprocess_run(["gh", "run", "download", "--repo", repo, "--name",
                              "pr", "--dir", "pr", str(run_id)])

        compare_folders(folder1, folder2, output_file)
    finally:
        if os.path.isdir(folder1):
            shutil.rmtree(folder1)

        if os.path.isdir(folder2):
            shutil.rmtree(folder2)

def get_comment_text(output_file, repo, run_id):
    size = os.path.getsize(output_file)
    if size == 0:
        print("No difference in the coverage reports")
        return

    comment = ":warning: The head of this PR and the base branch were compared for differences in the framework coverage reports. " + \
        "The generated reports are available in the [artifacts of this workflow run](https://github.com/" + repo + "/actions/runs/" + str(run_id) + "). " + \
        "The differences will be picked up by the nightly job after the PR gets merged. "

    if size < 2000:
        print("There's a small change in the CSV framework coverage reports")
        comment += "The following differences were found: \n\n"
        with open(output_file, 'r') as file:
            comment += file.read()
    else:
        print("There's a large change in the CSV framework coverage reports")
        comment += "The differences can be found in the " + \
            output_file + " artifact of this job."

    return comment

def comment_pr(output_file, repo, run_id):
    """
    Generates coverage diff produced by the changes in the current PR. If the diff is not empty, then post it as a comment.
    If a workflow run produces the same diff as the directly preceeding one, then don't post a comment.
    """

    # Store diff for current run
    write_diff_for_run(output_file, repo, run_id)

    try:
        with open("pr/NR") as file:
            pr_number = int(file.read())
    finally:
        if os.path.isdir("pr"):
            shutil.rmtree("pr")

    # Try storing diff for previous run:
    prev_output_file = "prev_" + output_file
    try:
        prev_run_id = get_previous_run_id(repo, run_id, pr_number)
        write_diff_for_run(prev_output_file, repo, prev_run_id)

        if filecmp.cmp(output_file, prev_output_file, shallow=False):
            print("Previous run " + str(prev_run_id) +
                  " resulted in the same diff, so not commenting again.")
            return
        else:
            print("Diff of previous run " +
                  str(prev_run_id) + " differs, commenting.")
    except Exception:
        # this is not mecessarily a failure, it can also mean that there was no previous run yet.
        print("Couldn't generate diff for previous run:", sys.exc_info()[1])
    finally:
        if os.path.isfile(prev_output_file):
            os.remove(prev_output_file)

    comment = get_comment_text(output_file, repo, run_id)
    post_comment(comment, repo, pr_number)

def post_comment(comment, repo, pr_number):
    print("Posting comment to PR #" + str(pr_number))
    utils.subprocess_run(["gh", "pr", "comment", str(pr_number),
                         "--repo", repo, "--body", comment])

def compare_folders(folder1, folder2, output_file):
    """
    Compares the contents of two folders and writes the differences to the output file.
    """

    languages = ['java']

    return_md = ""

    for lang in languages:
        expected_files = ""

        generated_output_rst = settings.generated_output_rst.format(
            language=lang)
        generated_output_csv = settings.generated_output_csv.format(
            language=lang)

        # check if files exist in both folder1 and folder 2
        if not check_file_exists(folder1 + "/" + generated_output_rst):
            expected_files += "- " + generated_output_rst + \
                " doesn't exist in folder " + folder1 + "\n"
        if not check_file_exists(folder2 + "/" + generated_output_rst):
            expected_files += "- " + generated_output_rst + \
                " doesn't exist in folder " + folder2 + "\n"
        if not check_file_exists(folder1 + "/" + generated_output_csv):
            expected_files += "- " + generated_output_csv + \
                " doesn't exist in folder " + folder1 + "\n"
        if not check_file_exists(folder2 + "/" + generated_output_csv):
            expected_files += "- " + generated_output_csv + \
                " doesn't exist in folder " + folder2 + "\n"

        if expected_files != "":
            print("Expected files are missing", file=sys.stderr)
            return_md += "\n### " + lang + "\n\n#### Expected files are missing for " + \
                lang + "\n" + expected_files + "\n"
            continue

        # compare contents of files
        cmp1 = compare_files_str(
            folder1 + "/" + generated_output_rst, folder2 + "/" + generated_output_rst)
        cmp2 = compare_files_str(
            folder1 + "/" + generated_output_csv, folder2 + "/" + generated_output_csv)

        if cmp1 != "" or cmp2 != "":
            print("Generated file contents are not matching", file=sys.stderr)
            return_md += "\n### " + lang + "\n\n#### Generated file changes for " + \
                lang + "\n\n"
            if cmp1 != "":
                return_md += "- Changes to " + generated_output_rst + \
                    ":\n```diff\n" + cmp1 + "```\n\n"
            if cmp2 != "":
                return_md += "- Changes to " + generated_output_csv + \
                    ":\n```diff\n" + cmp2 + "```\n\n"

    with open(output_file, 'w', newline='') as out:
        out.write(return_md)

def get_previous_run_id(repo, run_id, pr_number):
    """
    Gets the previous run id for a given workflow run, considering that the previous workflow run needs to come from the same PR.
    """

    # Get branch and repo from run:
    this_run = utils.subprocess_check_output(["gh", "api", "-X", "GET", "repos/" + repo + "/actions/runs/" + str(
        run_id), "--jq", "{ head_branch: .head_branch, head_repository: .head_repository.full_name }"])

    this_run = json.loads(this_run)
    pr_branch = this_run["head_branch"]
    pr_repo = this_run["head_repository"]

    # Get all previous runs that match branch, repo and workflow name:
    ids = utils.subprocess_check_output(["gh", "api", "-X", "GET", "repos/" + repo + "/actions/runs", "-f", "event=pull_request", "-f", "status=success", "-f", "name=\"" + artifacts_worflow_name + "\"", "--jq",
                                        "[.workflow_runs.[] | select(.head_branch==\"" + pr_branch + "\" and .head_repository.full_name==\"" + pr_repo + "\") | { created_at: .created_at, run_id: .id}] | sort_by(.created_at) | reverse | [.[].run_id]"])

    ids = json.loads(ids)
    if ids[0] != int(run_id):
        raise Exception("Expected to find " + str(run_id) +
                        " in the list of matching runs.")

    for previous_run_id in ids[1:]:
        utils.subprocess_run(["gh", "run", "download", "--repo", repo,
                             "--name", "pr", "--dir", "prev_run_pr", str(previous_run_id)])

        try:
            with open("prev_run_pr/NR") as file:
                prev_pr_number = int(file.read())
                print("PR number: " + str(prev_pr_number))
        finally:
            if os.path.isdir("prev_run_pr"):
                shutil.rmtree("prev_run_pr")

        # the previous run needs to be coming from the same PR:
        if pr_number == prev_pr_number:
            return previous_run_id

    raise Exception("Couldn't find previous run.")

output_file = sys.argv[1]
repo = sys.argv[2]
run_id = sys.argv[3]

comment_pr(output_file, repo, run_id)
<|endoftext|>


<commit_before>"""The decision module handles the three planning levels

Currently, only pre-specified planning is implemented.

The choices made in the three planning levels influence the set of interventions
and assets available within a model run.

The interventions available in a model run are stored in the
:class:`~smif.decision.intervention.InterventionRegister`.

When pre-specified planning are declared, each of the corresponding
interventions in the InterventionRegister are moved to the BuiltInterventionRegister.

Once pre-specified planning is instantiated, the action space for rule-based and
optimisation approaches can be generated from the remaining Interventions in the
InterventionRegister.

"""
__author__ = "Will Usher, Tom Russell"
__copyright__ = "Will Usher, Tom Russell"
__license__ = "mit"

from abc import ABCMeta, abstractmethod
from logging import getLogger

from smif.decision.intervention import InterventionRegister

class DecisionManager(object):
    """A DecisionManager is initialised with one or more model run strategies that refer to
    DecisionModules such as pre-specified planning, a rule-based models or multi-objective
    optimisation. These implementations influence the combination and ordering of decision
    iterations and model timesteps that need to be performed by the model runner.

    The DecisionManager presents a simple decision loop interface to the model runner, in the
    form of a generator which allows the model runner to iterate over the collection of
    independent simulations required at each step.

    (Not yet implemented.) A DecisionManager collates the output of the decision algorithms and
    writes the post-decision state through a DataHandle. This allows Models to access a given
    decision state (identified uniquely by timestep and decision iteration id).

    (Not yet implemented.) A DecisionManager may also pass a DataHandle down to a
    DecisionModule, allowing the DecisionModule to access model results from previous timesteps
    and decision iterations when making decisions.

    Arguments
    ---------
    timesteps: list
    strategies: list
    """

    def __init__(self, timesteps, strategies, interventions):

        self.logger = getLogger(__name__)

        self._timesteps = timesteps
        self._strategies = strategies
        self._decision_modules = []

        self._set_up_decision_modules()

        self.register = InterventionRegister()
        for intervention in interventions:
            self.register.register(intervention)

    def _set_up_decision_modules(self):

        self.logger.info("%s strategies found", len(self._strategies))
        interventions = []

        for index, strategy in enumerate(self._strategies):
            if strategy['strategy'] == 'pre-specified-planning':

                msg = "Adding %s interventions to pre-specified-planning %s"
                self.logger.info(msg, len(strategy['interventions']), index)

                interventions.extend(strategy['interventions'])

            else:
                msg = "Only pre-specified planning strategies are implemented"
                raise NotImplementedError(msg)

        if interventions:
            self._decision_modules.append(
                PreSpecified(self._timesteps, interventions)
                )

    def decision_loop(self):
        """Generate bundles of simulation steps to run.

        Each iteration returns a dict: {decision_iteration (int) => list of timesteps (int)}

        With only pre-specified planning, there is a single step in the loop, with a single
        decision iteration with timesteps covering the entire model horizon.

        With a rule based approach, there might be many steps in the loop, each with a single
        decision iteration and single timestep, moving on once some threshold is satisfied.

        With a genetic algorithm, there might be a configurable number of steps in the loop,
        each with multiple decision iterations (one for each member of the algorithm's
        population) and timesteps covering the entire model horizon.

        Implicitly, if the bundle returned in an iteration contains multiple decision
        iterations, they can be performed in parallel. If each decision iteration contains
        multiple timesteps, they can also be parallelised, so long as there are no temporal
        dependencies.
        """
        if len(self._decision_modules) > 0:
            for module in self._decision_modules:
                yield module._get_next_decision_iteration()
        else:
            yield {0: [x for x in self._timesteps]}

    def get_decision(self, timestep, iteration):
        """Return all interventions built in the given timestep

        for the given decision
        iteration.

        Arguments
        ---------
        timestep : int
            A timestep (planning year)
        iteration : int
            A decision iteration

        """
        decisions = []
        for module in self._decision_modules:
            decisions.extend(module.get_decision(timestep, iteration))
        self.logger.debug(
            "Retrieved %s decisions from %s",
            len(decisions), str(self._decision_modules))
        return decisions

class DecisionModule(metaclass=ABCMeta):
    """Abstract class which provides the interface to decision mechanisms.

    These mechanisms including Pre-Specified Planning, a Rule-based Approach and
    Multi-objective Optimisation.

    This class provides two main public methods, ``__next__`` which is normally
    called implicitly as a call to the class as an iterator, and ``get_decision()``
    which takes as arguments a smif.model.Model object, and ``timestep`` and
    ``decision_iteration`` integers. The first of these returns a dict of
    decision_iterations and timesteps over which a SosModel should be iterated.
    The latter provides a means to furnish the structure of contained Model
    objects through a list of historical and recent interventions.

    Arguments
    ---------
    timesteps : list
        A list of planning timesteps

    """
    def __init__(self, timesteps):
        self.timesteps = timesteps

    def __next__(self):
        return self._get_next_decision_iteration()

    @abstractmethod
    def _get_next_decision_iteration(self):
        """Implement to return the next decision iteration

        Within a list of decision-iteration/timestep pairs, the assumption is
        that all decision iterations can be run in parallel
        (otherwise only one will be returned) and within a decision interation,
        all timesteps may be run in parallel as long as there are no
        inter-timestep state transitions required (e.g. reservoir level)

        Returns
        -------
        dict
            Yields a dictionary keyed by decision iteration (int) whose values contain
            a list of timesteps
        """
        raise NotImplementedError

    @abstractmethod
    def get_decision(self, timestep, iteration):
        """Return decisions for a given timestep and decision iteration
        """
        raise NotImplementedError

    def _get_previous_state(self, timestep, decision_iteration):
        """Gets state of the previous `timestep` for `decision_iteration`

        Arguments
        ---------
        timestep : int
        decision_iteration : int

        Returns
        -------
        numpy.ndarray
        """
        return self.get_decision(timestep.PREVIOUS, decision_iteration)

    def _set_post_decision_state(self, timestep, decision_iteration, decision):
        """Sets the post-decision state

        Arguments
        ---------
        timestep : int
        decision_iteration : int
        decision : numpy.ndarray

        Notes
        -----
        `decision` should contain only the newly decided interventions
        """
        state = self._get_previous_state(timestep, decision_iteration)
        post_decision_state = state.bitwise_or(decision)
        self.register.set_state(timestep, decision_iteration, post_decision_state)

    @abstractmethod
    def _set_state(self, timestep, decision_iteration):
        """Implement to set the current state of a Model

        Arguments
        ---------
        timestep : int
        decision_iteration : int

        Notes
        -----
        1. Get state at previous timestep
        2. Compute decisions (interventions at current timestep)
        3. Accumulate to create current state
        4. Write this via the DataHandle to DataInterface

        This is a candidate for memoization

        """
        raise NotImplementedError

class PreSpecified(DecisionModule):
    """Pre-specified planning

    Arguments
    ---------
    timesteps : list
    planned_interventions : list
        A list of dicts ``{'name': 'intervention_name', 'build_year': 2010}``
        representing historical or planned interventions
    """

    def __init__(self, timesteps, planned_interventions):
        super().__init__(timesteps)

        self._planned = planned_interventions

    def _get_next_decision_iteration(self):
        return {0: [x for x in self.timesteps]}

    def _set_state(self, timestep, decision_iteration):
        """Pre-specified planning interventions are loaded during initialisation

        Pre-specified planning interventions are loaded during initialisation
        of a model run. This method just needs to copy the existing system state
        to the correct decision iteration reference.

        Arguments
        ---------
        timestep : int
            A timestep (planning year)
        iteration : int
            A decision iteration
        """
        pass

    def get_decision(self, timestep, iteration=None):
        """Return a dict of intervention names built in timestep

        Arguments
        ---------
        timestep : int
            A timestep (planning year)
        iteration : int
            A decision iteration

        Returns
        -------
        list of tuples

        Examples
        --------
        >>> dm = PreSpecified([2010, 2015], [{'name': 'intervention_a', 'build_year': 2010}])
        >>> dm.get_decision(2010)
        [{'name': intervention_a', 'build_year': 2010}]
        """
        decisions = []

        assert isinstance(self._planned, list)

        for intervention in self._planned:
            build_year = int(intervention['build_year'])
            if self.buildable(build_year, timestep):
                decisions.append(intervention)
        return decisions

    def buildable(self, build_year, timestep):
        """Interventions are deemed available if build_year is less than next timestep

        For example, if `a` is built in 2011 and timesteps are
        [2005, 2010, 2015, 2020] then buildable returns True for timesteps
        2010, 2015 and 2020 and False for 2005.
        """
        if not isinstance(build_year, (int, float)):
            msg = "Build Year should be an integer but is a {}"
            raise TypeError(msg.format(type(build_year)))
        if timestep not in self.timesteps:
            raise ValueError("Timestep not in model timesteps")
        index = self.timesteps.index(timestep)
        if index == len(self.timesteps) - 1:
            next_year = timestep + 1
        else:
            next_year = self.timesteps[index + 1]

        if int(build_year) < next_year:
            return True
        else:
            return False

class RuleBased(DecisionModule):
    """Rule-base decision modules
    """

    def __init__(self, timesteps):
        super().__init__(timesteps)
        self.satisfied = False
        self.current_timestep_index = 0
        self.current_iteration = 0

    def _get_next_decision_iteration(self):
            if self.satisfied and self.current_timestep_index == len(self.timesteps) - 1:
                return None
            elif self.satisfied and self.current_timestep_index <= len(self.timesteps):
                self.satisfied = False
                self.current_timestep_index += 1
                self.current_iteration += 1
                return {self.current_iteration: [self.timesteps[self.current_timestep_index]]}
            else:
                self.current_iteration += 1
                return {self.current_iteration: [self.timesteps[self.current_timestep_index]]}

    def get_decision(self, timestep, iteration):
        return []

    def _set_state(self, timestep, decision_iteration):
        raise NotImplementedError
<commit_msg>Pass InterventionRegister to DecisionModules<commit_after>"""The decision module handles the three planning levels

Currently, only pre-specified planning is implemented.

The choices made in the three planning levels influence the set of interventions
and assets available within a model run.

The interventions available in a model run are stored in the
:class:`~smif.decision.intervention.InterventionRegister`.

When pre-specified planning are declared, each of the corresponding
interventions in the InterventionRegister are moved to the BuiltInterventionRegister.

Once pre-specified planning is instantiated, the action space for rule-based and
optimisation approaches can be generated from the remaining Interventions in the
InterventionRegister.

"""
__author__ = "Will Usher, Tom Russell"
__copyright__ = "Will Usher, Tom Russell"
__license__ = "mit"

from abc import ABCMeta, abstractmethod
from logging import getLogger

from smif.decision.intervention import InterventionRegister

class DecisionManager(object):
    """A DecisionManager is initialised with one or more model run strategies that refer to
    DecisionModules such as pre-specified planning, a rule-based models or multi-objective
    optimisation. These implementations influence the combination and ordering of decision
    iterations and model timesteps that need to be performed by the model runner.

    The DecisionManager presents a simple decision loop interface to the model runner, in the
    form of a generator which allows the model runner to iterate over the collection of
    independent simulations required at each step.

    (Not yet implemented.) A DecisionManager collates the output of the decision algorithms and
    writes the post-decision state through a DataHandle. This allows Models to access a given
    decision state (identified uniquely by timestep and decision iteration id).

    (Not yet implemented.) A DecisionManager may also pass a DataHandle down to a
    DecisionModule, allowing the DecisionModule to access model results from previous timesteps
    and decision iterations when making decisions.

    Arguments
    ---------
    timesteps: list
    strategies: list
    """

    def __init__(self, timesteps, strategies, interventions):

        self.logger = getLogger(__name__)

        self._timesteps = timesteps
        self._strategies = strategies
        self._decision_modules = []

        self._set_up_decision_modules()

        self.register = InterventionRegister()
        for intervention in interventions:
            self.register.register(intervention)

    def _set_up_decision_modules(self):

        self.logger.info("%s strategies found", len(self._strategies))
        interventions = []

        for index, strategy in enumerate(self._strategies):
            if strategy['strategy'] == 'pre-specified-planning':

                msg = "Adding %s interventions to pre-specified-planning %s"
                self.logger.info(msg, len(strategy['interventions']), index)

                interventions.extend(strategy['interventions'])

            else:
                msg = "Only pre-specified planning strategies are implemented"
                raise NotImplementedError(msg)

        if interventions:
            self._decision_modules.append(
                PreSpecified(self._timesteps, interventions, self.register)
                )

    def decision_loop(self):
        """Generate bundles of simulation steps to run.

        Each iteration returns a dict: {decision_iteration (int) => list of timesteps (int)}

        With only pre-specified planning, there is a single step in the loop, with a single
        decision iteration with timesteps covering the entire model horizon.

        With a rule based approach, there might be many steps in the loop, each with a single
        decision iteration and single timestep, moving on once some threshold is satisfied.

        With a genetic algorithm, there might be a configurable number of steps in the loop,
        each with multiple decision iterations (one for each member of the algorithm's
        population) and timesteps covering the entire model horizon.

        Implicitly, if the bundle returned in an iteration contains multiple decision
        iterations, they can be performed in parallel. If each decision iteration contains
        multiple timesteps, they can also be parallelised, so long as there are no temporal
        dependencies.
        """
        if len(self._decision_modules) > 0:
            for module in self._decision_modules:
                yield module._get_next_decision_iteration()
        else:
            yield {0: [x for x in self._timesteps]}

    def get_decision(self, timestep, iteration):
        """Return all interventions built in the given timestep

        for the given decision
        iteration.

        Arguments
        ---------
        timestep : int
            A timestep (planning year)
        iteration : int
            A decision iteration

        """
        decisions = []
        for module in self._decision_modules:
            decisions.extend(module.get_decision(timestep, iteration))
        self.logger.debug(
            "Retrieved %s decisions from %s",
            len(decisions), str(self._decision_modules))
        return decisions

class DecisionModule(metaclass=ABCMeta):
    """Abstract class which provides the interface to decision mechanisms.

    These mechanisms including Pre-Specified Planning, a Rule-based Approach and
    Multi-objective Optimisation.

    This class provides two main public methods, ``__next__`` which is normally
    called implicitly as a call to the class as an iterator, and ``get_decision()``
    which takes as arguments a smif.model.Model object, and ``timestep`` and
    ``decision_iteration`` integers. The first of these returns a dict of
    decision_iterations and timesteps over which a SosModel should be iterated.
    The latter provides a means to furnish the structure of contained Model
    objects through a list of historical and recent interventions.

    Arguments
    ---------
    timesteps : list
        A list of planning timesteps
    intervention_register : smif.interventions.InterventionRegister
        Reference to a fully populated intervention register

    """
    def __init__(self, timesteps, intervention_register):
        self.timesteps = timesteps
        self.intervention_register = intervention_register

    def __next__(self):
        return self._get_next_decision_iteration()

    @abstractmethod
    def _get_next_decision_iteration(self):
        """Implement to return the next decision iteration

        Within a list of decision-iteration/timestep pairs, the assumption is
        that all decision iterations can be run in parallel
        (otherwise only one will be returned) and within a decision interation,
        all timesteps may be run in parallel as long as there are no
        inter-timestep state transitions required (e.g. reservoir level)

        Returns
        -------
        dict
            Yields a dictionary keyed by decision iteration (int) whose values contain
            a list of timesteps
        """
        raise NotImplementedError

    @abstractmethod
    def get_decision(self, timestep, iteration):
        """Return decisions for a given timestep and decision iteration
        """
        raise NotImplementedError

    def _get_previous_state(self, timestep, decision_iteration):
        """Gets state of the previous `timestep` for `decision_iteration`

        Arguments
        ---------
        timestep : int
        decision_iteration : int

        Returns
        -------
        numpy.ndarray
        """
        return self.get_decision(timestep.PREVIOUS, decision_iteration)

    def _set_post_decision_state(self, timestep, decision_iteration, decision):
        """Sets the post-decision state

        Arguments
        ---------
        timestep : int
        decision_iteration : int
        decision : numpy.ndarray

        Notes
        -----
        `decision` should contain only the newly decided interventions
        """
        state = self._get_previous_state(timestep, decision_iteration)
        post_decision_state = state.bitwise_or(decision)
        self.register.set_state(timestep, decision_iteration, post_decision_state)

    @abstractmethod
    def _set_state(self, timestep, decision_iteration):
        """Implement to set the current state of a Model

        Arguments
        ---------
        timestep : int
        decision_iteration : int

        Notes
        -----
        1. Get state at previous timestep
        2. Compute decisions (interventions at current timestep)
        3. Accumulate to create current state
        4. Write this via the DataHandle to DataInterface

        This is a candidate for memoization

        """
        raise NotImplementedError

class PreSpecified(DecisionModule):
    """Pre-specified planning

    Arguments
    ---------
    timesteps : list
    planned_interventions : list
        A list of dicts ``{'name': 'intervention_name', 'build_year': 2010}``
        representing historical or planned interventions
    """

    def __init__(self, timesteps, intervention_register, planned_interventions):
        super().__init__(timesteps, intervention_register)
        self._planned = planned_interventions

    def _get_next_decision_iteration(self):
        return {0: [x for x in self.timesteps]}

    def _set_state(self, timestep, decision_iteration):
        """Pre-specified planning interventions are loaded during initialisation

        Pre-specified planning interventions are loaded during initialisation
        of a model run. This method just needs to copy the existing system state
        to the correct decision iteration reference.

        Arguments
        ---------
        timestep : int
            A timestep (planning year)
        iteration : int
            A decision iteration
        """
        pass

    def get_decision(self, timestep, iteration=None):
        """Return a dict of intervention names built in timestep

        Arguments
        ---------
        timestep : int
            A timestep (planning year)
        iteration : int
            A decision iteration

        Returns
        -------
        list of tuples

        Examples
        --------
        >>> dm = PreSpecified([2010, 2015], [{'name': 'intervention_a', 'build_year': 2010}])
        >>> dm.get_decision(2010)
        [{'name': intervention_a', 'build_year': 2010}]
        """
        decisions = []

        assert isinstance(self._planned, list)

        for intervention in self._planned:
            build_year = int(intervention['build_year'])

            data = self.intervention_register.get_intervention(intervention['name'])
            lifetime = data['lifetime']

            if self.buildable(build_year, timestep) and \
               self.within_lifetime(build_year, timestep, lifetime):
                decisions.append(intervention)
        return decisions

    def buildable(self, build_year, timestep):
        """Interventions are deemed available if build_year is less than next timestep

        For example, if `a` is built in 2011 and timesteps are
        [2005, 2010, 2015, 2020] then buildable returns True for timesteps
        2010, 2015 and 2020 and False for 2005.
        """
        if not isinstance(build_year, (int, float)):
            msg = "Build Year should be an integer but is a {}"
            raise TypeError(msg.format(type(build_year)))
        if timestep not in self.timesteps:
            raise ValueError("Timestep not in model timesteps")
        index = self.timesteps.index(timestep)
        if index == len(self.timesteps) - 1:
            next_year = timestep + 1
        else:
            next_year = self.timesteps[index + 1]

        if int(build_year) < next_year:
            return True
        else:
            return False

    def within_lifetime(self, build_year, timestep, lifetime):
        """Interventions are deemed active if build_year + lifetime <= timestep

        Arguments
        ---------
        build_year : int
        timestep : int
        lifetime : int
        """
        if build_year + lifetime <= timestep:
            return True
        else:
            return False

class RuleBased(DecisionModule):
    """Rule-base decision modules
    """

    def __init__(self, timesteps):
        super().__init__(timesteps)
        self.satisfied = False
        self.current_timestep_index = 0
        self.current_iteration = 0

    def _get_next_decision_iteration(self):
            if self.satisfied and self.current_timestep_index == len(self.timesteps) - 1:
                return None
            elif self.satisfied and self.current_timestep_index <= len(self.timesteps):
                self.satisfied = False
                self.current_timestep_index += 1
                self.current_iteration += 1
                return {self.current_iteration: [self.timesteps[self.current_timestep_index]]}
            else:
                self.current_iteration += 1
                return {self.current_iteration: [self.timesteps[self.current_timestep_index]]}

    def get_decision(self, timestep, iteration):
        return []

    def _set_state(self, timestep, decision_iteration):
        raise NotImplementedError
<|endoftext|>