Environment.py

import string

class Env:
    nodeToNamespace = dict({
        'FieldNode': 'fieldDcl',
        'InterNode': 'type',
        'ClassNode': 'type',
        'MethodNode': 'method',
        'VarDclNode': 'expr'

    })

    def __init__(self, parentEnv):
        self.parentEnv = parentEnv
        self.map = {}
                #  entries in the map:
                # type: (typeName, 'type') -> node
                # method: (methodName. 'method') -> {paramTypes -> node} for method overloding
                # expr: (name, 'expr') -> node

    def addtoEnv(self, node):
        # use node.__class__.__name__ to get namespace
        nodeType = node.__class__.__name__
        if nodeType not in self.nodeToNamespace:
            namespace = 'expr'
        else:
            namespace = self.nodeToNamespace.get(nodeType)

        # Add to Env
        key = (node.name, namespace)
        # add entry in nested map for method overloding
        if namespace == 'method':
            if key not in self.map:
                self.map[key] = {}
            self.map[key][node.paramTypes] = node
        else:
            if key in self.map:
                raise Exception('ERROR: Declaration of {} is already in current Environment'.format(node.name))
            if namespace == 'expr':
                self.map[key] = (False, node)
            else:
                self.map[key] = node


    def getNode(self, name, namespace):
        key = (name, namespace)

        if key in self.map:
            return self.map.get(key)
        elif self.parentEnv:
            return self.parentEnv.getNode(name, namespace)
        raise Exception("ERROR: Can't find definition of {} in the Environment".format(name))

    # A wrapper around getNode to find if node exists in environment already 
    def findNode(self, name, namespace):
        try: 
            self.getNode(name, namespace) 
        except:
            return False # node is not found in environment 
        
        return True 


###################################

class GlobalEnv(Env):
    def __init__(self):
        super().__init__(None)
        #self.map = {} # {Canonical name -> node}
        self.packageMap = {} # {packageName -> {(simpleName, namespace) -> node}}
                            #   A map of maps
                            #   contains duplicate info as self.map for easy searching

    # input: Compiliation Unit Node (CompNode)
    def addtoEnv(self, node):
        pName = node.packageName
        typeDcl = node.typeDcl
        typeDclType = typeDcl.__class__.__name__
        # at global level there are only class and interface
        namespace = self.nodeToNamespace.get(typeDclType)

        # Add to map
        mapKey = pName + '.'  + typeDcl.name
        if mapKey in self.map:
            raise Exception('ERROR: Declaration of {} is already in current Environment'.format(pName + '.'  + typeDcl.name))
        self.map[mapKey] = typeDcl

        # Add to packageMap
        pMapKey = (typeDcl.name, namespace)
        if not pName in self.packageMap:
            self.packageMap[pName] = {}

        if pMapKey in self.packageMap[pName]:
            raise Exception('ERROR: Declaration of {} is already in current Environment'.format(pName + '.'  + typeDcl.name))
        self.packageMap[pName][pMapKey] = typeDcl

    # Use getNode() from base class to get node using Canonical Name (full name)
    def getNode(self, key, imported, packageName):
        # TODO: error if ambiguous
        # TODO: not sure whether name is canonical or not????
        #  i.e. there are more than one of this type. Currently it only gets the first occurence
        name = key[0]
        
        # 1. enclosing class/interface
        #    - already did
        # 2. single-type import
        if name in imported and name in self.map:
            return self.map.get(name)
        # 3. type in the current package
        full = packageName + '.' + name
        # print(full)
        if full in self.map:
            return self.map.get(full)
        # 4. import-on-demand
        for i in imported:
            if '*' in i:
                full = i.replace("*", name)
                if full in self.map:
                    return self.map.get(full)
        # imported implicitly: java.io.*, java.lang.*, java.util.*
        implicitly = ['java.lang.', 'java.io.', 'java.util.']
        for i in implicitly:
            if i in name and name in self.map:
                return self.map.get(name)
            elif (i+name) in self.map:
                return self.map.get(i+name)

        raise Exception("ERROR: Can't find definition of {} in the Environment".format(key))

    # method for getting all the nodes under a package (import All)
    # returns a dict of types under that package name
    def getNodesByPackage(self, pName):
        if not pName in self.packageMap:
            raise Exception("ERROR: Can't find definition of package {} in the Environment".format(pName))
        return self.packageMap.get(pName)

###################################

class CompEnv(Env):
    def __init__(self, parentEnv, imported, packageName):
        super().__init__(parentEnv) # globalEnv is parentEnv
        self.imported = imported # list of strings that have been imported in this comp
        self.packageName = packageName # string of the current file's package name

    def getNode(self, name, namespace):
        key = (name, namespace)

        if key in self.map:
            return self.map.get(key)
        elif self.parentEnv: # globalEnv
            return self.parentEnv.getNode(key, self.imported, self.packageName)
        raise Exception("ERROR: Can't find definition of {} in the Environment".format(name))