Cleaned up some files

Author: willgebhardt

ngcsimlib/_src/compartment/compartment.py

@@ -2,12 +2,13 @@
 from ngcsimlib._src.global_state.manager import global_state_manager as gState
 from ngcsimlib._src.logger import warn
 import ast
-from typing import TypeVar, Union
+from typing import TypeVar, Union, Set
 from ngcsimlib._src.operations.BaseOp import BaseOp
 from ngcsimlib._src.context.context_manager import global_context_manager as gcm
 
 T = TypeVar('T')
 
+
 class Compartment(metaclass=CompartmentMeta):
     """
     Compartments exist as a layer between the global state of models and the
@@ -19,17 +20,29 @@ class Compartment(metaclass=CompartmentMeta):
     nothing will access it unless the user manually goes looking for it). The
     compartment should be reflected in the global state immediately after it is
     initialized. Compartments can be flagged as fixed which means that they
-    exist in the global state and can be used in compiled methods but they can
+    exist in the global state and can be used in compiled methods, but they can
     not be changed after creation.
 
     Args
         initial_value: the initial value to set in the global state
 
         fixed (default=False): sets the flag for if this compartment is fixed.
+
+        display_name (default=None): sets the display name of the compartment
+
+        units (default=None): sets the units of the compartment
+
+        plot_method (default=None): sets the plot method of the compartment,
+            this method is to be used by the processes when monitoring this
+            compartment to integrate with the plotting system.
     """
-    def __init__(self, initial_value: T, fixed: bool = False,
-                 display_name=None, units=None, plot_method=None):
-        self._initial_value = initial_value
+    def __init__(self, initial_value: T,
+                 fixed: bool = False,
+                 display_name: str | None = None,
+                 units: str | None = None,
+                 plot_method: str | None = None):
+
+        self._initial_value: T = initial_value
 
         self.name = None
         self._root_target = None
@@ -41,7 +54,7 @@ def __init__(self, initial_value: T, fixed: bool = False,
         self.plot_method = plot_method
 
     @property
-    def root(self):
+    def root(self) -> str | None:
         return self._root_target
 
     @property
@@ -91,7 +104,11 @@ def get(self) -> T:
         """
         return self._get_value()
 
-    def get_needed_keys(self):
+    def get_needed_keys(self) -> Set[str]:
+        """
+        Returns: Returns a set of compartment paths that are needed to compute
+            the value of this compartment
+        """
         if isinstance(self.target, BaseOp):
             return self.target.get_needed_keys()
         return set(self.target)
@@ -130,7 +147,7 @@ def __rshift__(self, other):
             other.__rrshift__(self)
 
     @property
-    def target(self):
+    def target(self) -> Union["BaseOp", str]:
         """
         Returns: the current target of the compartment
         """
@@ -156,4 +173,4 @@ def target(self, value: Union["Compartment", "BaseOp", str]):
         if isinstance(value, str):
             self._target = value
 
-        raise ValueError("Invalid compartment target ", value)
+        raise ValueError("Invalid compartment target ", value)

ngcsimlib/_src/context/context.py

@@ -40,6 +40,7 @@ class Context(object):
     with block. This means that in order to use any of the compiled methods or
     processes defined the with block must first be left.
     """
+
     def __new__(cls, name: str, *args, **kwargs):
         targetPath = gcm.append_path(addition=name)
         if gcm.exists(targetPath):
@@ -59,8 +60,7 @@ def __init__(self, name: str):
 
         self.name = name
         self.objects = {}
-        self._connections = {}
-
+        self._connections: Dict[str: Union["Compartment", "BaseOp"]] = {}
 
     def __enter__(self):
         self.__previous_path = gcm.current_path
@@ -72,7 +72,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         gcm.step_to(self.__previous_path)
         self.__previous_path = None
 
-    def recompile(self):
+    def recompile(self) -> None:
         """
         Recompiles all the context aware objects inside the context based on
         their priority. The higher the priority is, the sooner it will happen.
@@ -86,38 +86,39 @@ def recompile(self):
         """
         priorities = {}
 
-
         for objectType in self.objects.keys():
             _objs = self.get_objects_by_type(objectType)
             for objName, obj in _objs.items():
                 if getattr(obj, "_is_compilable", False):
-                    p = getattr(obj, "_priority", None)
-                    p = 0 if p is None else p
+                    p = getattr(obj, "_priority", None) or 0
 
                     if p not in priorities:
                         priorities[p] = []
 
                     priorities[p].append(obj)
 
-
         keys = sorted(priorities.keys(), reverse=True)
         for key in keys:
             for obj in priorities[key]:
                 obj.compile()
 
-    def registerObj(self, obj: "ContextAwareObjectMeta"):
+    def registerObj(self, obj: "ContextAwareObjectMeta") -> bool:
         """
         Registers an object in the context. The context automatically sorts the
-        objects by type through the "_type" field set on the object/class. It
-        expects to be given a "ContextObjectType" but it is not a requirement.
+        objects by type through the "_type" field set on the object/class.
+        Standard practice is to use the predefined decorators or superclasses
+        found in this library to set this field, but it is not a requirement.
         If an unknown type is provided to the context it will still sort it
         into a bin with other objects of the same type. (Note: _type can be
         either a string or ContextObjectTypes.TYPE, both will be grouped
         together)
 
         Args:
-            obj: The object to register, requires the "_type" field to be
-                defined and not null
+            obj: The object to register in the context
+
+        Returns:
+            boolean: marks if the object was successfully registered in the
+                context
         """
         _type = getattr(obj, "_type", None)
         if _type is None:
@@ -126,10 +127,12 @@ def registerObj(self, obj: "ContextAwareObjectMeta"):
                  f"object will be limited. Please use one of the provided "
                  f"context object types or define your own to ensure "
                  f"compatability")
-            return
+            return False
 
-        if not isinstance(_type, ContextObjectTypes) and not (
-            isinstance(_type, str) and _type in ContextObjectTypes.__members__):
+        if (not isinstance(_type, ContextObjectTypes) and
+            _type not in self.objects.keys() and
+            not (isinstance(_type, str) and _type in ContextObjectTypes.__members__)
+            ):
             warn(
                 f"Context object type {_type} is not known to this context. It will "
                 f"be stored and tracked but some functionality will be "
@@ -146,9 +149,10 @@ def registerObj(self, obj: "ContextAwareObjectMeta"):
             warn(f"Trying to register context object with the same name "
                  f"({obj.name}) as another object in this context. Aborting "
                  f"registration!")
-            return
+            return False
 
         self.objects[_type][obj.name] = obj
+        return True
 
     def get_objects_by_type(self, objectType: ContextObjectTypes | str) -> Dict[
         str, "ContextAwareObjectMeta"]:
@@ -213,20 +217,18 @@ def get_objects(self, *object_names: str,
             return _objs[0]
         return _objs
 
-
     def get_components(self, *component_names: str, unwrap: bool = True) -> \
         Union[None, "ContextAwareObjectMeta", List[Union[
             "ContextAwareObjectMeta", None]]]:
         return self.get_objects(*component_names,
                                 objectType=ContextObjectTypes.component,
                                 unwrap=unwrap)
 
-    def add_connection(self, source, destination):
+    def add_connection(self, source: Union["Compartment", "BaseOp"], destination: "Compartment"):
         self._connections[destination.root] = source
 
-
     def save_to_json(self, directory: str, model_name: Union[str, None] = None,
-                     custom_save: bool = True, overwrite: bool = False):
+                     custom_save: bool = True, overwrite: bool = False) -> None:
         """
         Saves the context to a collection fo JSON files.
 
@@ -256,7 +258,6 @@ def save_to_json(self, directory: str, model_name: Union[str, None] = None,
                     print('Failed to delete %s. Reason: %s' % (file_path, e))
             shutil.rmtree(directory + "/" + model_name)
 
-
         path = make_unique_path(directory, model_name)
 
         contextMeta = {"types": list(self.objects.keys()),
@@ -276,7 +277,8 @@ def save_to_json(self, directory: str, model_name: Union[str, None] = None,
 
             for obj_name, obj in _objs.items():
                 objData = {}
-                if hasattr(obj, "to_json") and callable(getattr(obj, "to_json")):
+                if hasattr(obj, "to_json") and callable(
+                    getattr(obj, "to_json")):
                     objData.update(obj.to_json())
 
                 objData["modulePath"] = modManager.resolve_public_import(obj)
@@ -303,10 +305,8 @@ def save_to_json(self, directory: str, model_name: Union[str, None] = None,
         with open(f"{path}/connections.json", "w") as fp:
             json.dump(connections, fp, indent=4)
 
-
-
     @staticmethod
-    def load(directory: str, module_name: str):
+    def load(directory: str, module_name: str) -> "Context":
         if gcm.exists(gcm.append_path(module_name)):
             warn("Trying to load a context that already exists, returning "
                  "existing context")
@@ -330,17 +330,20 @@ def load(directory: str, module_name: str):
                     kwargs = objData["kwargs"]
                     newObj = objKlass(*args, **kwargs)
 
-                    delayed_load.append((getattr(newObj, "_priority", 0), newObj, objData, type_path))
+                    delayed_load.append((
+                        getattr(newObj, "_priority", 0), newObj,
+                        objData, type_path))
 
-            delayed_load = sorted(delayed_load, key=lambda x: x[0], reverse=True)
+            delayed_load = sorted(delayed_load, key=lambda x: x[0],
+                                  reverse=True)
             for _, obj, data, type_path in delayed_load:
-                if hasattr(obj, "from_json") and callable(getattr(obj, "from_json")):
+                if hasattr(obj, "from_json") and callable(
+                    getattr(obj, "from_json")):
                     obj.from_json(objData)
 
                 if hasattr(obj, "load") and callable(getattr(obj, "load")):
                     obj.load(f"{type_path}/custom")
 
-
             with open(f"{path}/connections.json", "r") as fp:
                 connectionData = json.load(fp)
                 for connectionRoot, target in connectionData.items():
@@ -350,8 +353,4 @@ def load(directory: str, module_name: str):
                     else:
                         dest.target = BaseOp.load_op(target)
 
-
         return ctx
-
-
-

ngcsimlib/_src/context/contextAwareObject.py

@@ -44,7 +44,7 @@ def to_json(self) -> Dict[str, Any]:
                 "kwargs": safe_kwargs}
         return data
 
-    def compile(self):
+    def compile(self) -> None:
         """
         A wrapper to compile this object, unless a custom compiler is being used
         do not modify this method.

ngcsimlib/_src/context/contextAwareObjectMeta.py

@@ -4,6 +4,7 @@
 from .context_manager import global_context_manager as gcm
 from collections.abc import Iterable
 
+
 def extract_name(cls, args, kwargs):
     init = cls.__init__
     sig = inspect.signature(init)
@@ -14,6 +15,7 @@ def extract_name(cls, args, kwargs):
         return bound.arguments["name"]
     return None
 
+
 class ContextAwareObjectMeta(type):
     def __new__(cls, name, bases, attrs):
         if '__enter__' not in attrs:

ngcsimlib/_src/context/contextObjectDecorators.py

@@ -15,13 +15,6 @@ def process(cls):
         cls._type = ContextObjectTypes.process
         return cls
 
-    # @staticmethod
-    # def _operation(cls):
-    #     cls._type = ContextObjectTypes.operation
-    #     return cls
-
-
 
 component = ContextObjectDecorators.component
 process = ContextObjectDecorators.process
-# operation = ContextObjectDecorators._operation

ngcsimlib/_src/parser/contextTransformer.py

@@ -106,7 +106,7 @@ def visit_Call(self, node):
             subAttr = getattr(attr, node.func.attr)
             if not hasattr(subAttr, "compiled"):
                 error("Attempting to use a method of a subcomponent that is not compiled/compilable")
-
+            
             method_id = f"{attr.context_path.replace(':', '_')}_{node.func.attr}"
             subAst = subAttr.compiled.ast
             subAst.body[0].body = subAst.body[0].body[:-1]

ngcsimlib/_src/parser/utils.py

@@ -20,7 +20,7 @@ class CompiledMethod:
     def __init__(self, fn, fn_ast, auxiliary_ast, namespace, extra_globals):
         self._fn = fn
         self._fn_ast = fn_ast
-        self._auxiliary_ast = auxiliary_ast if auxiliary_ast is not None else {}
+        self._auxiliary_ast = auxiliary_ast or {}
         self._namespace = namespace
         self._extra_globals = extra_globals