Remove classes for old import syntax, and convert between both

src/gboml/ast/__init__.py

@@ -4,7 +4,7 @@ __all__ = [
     "VarOrParam", "Array", "Loop", "BaseLoop", "Function", "BoolExpressionOp",
     "BoolExpressionComparison", "ScopeChange", "ImportFile", "Definition", "Constraint",
     "StdConstraint", "SOSConstraint", "Objective", "VariableDefinition", "Node",
-    "HyperEdge", "NodeDefinition", "NodeImport", "HyperEdgeDefinition", "HyperEdgeImport",
+    "HyperEdge", "NodeDefinition", "HyperEdgeDefinition",
     "ExpressionOp", "GBOMLGraph", "ImplicitLoop", "RValue", "RValueWithGen", "GeneratedRValue",
     "Range", "MultiLoop", "DictEntry", "Dictionary", "NodeGenerator", "HyperEdgeGenerator",
     "DefinitionType", "FunctionDefinition", "ConstantDefinition", "ExpressionDefinition",

src/gboml/ast/hyperedges.py

@@ -33,11 +33,3 @@ class HyperEdgeGenerator(HyperEdge):
     constraints: list[Constraint] = field(default_factory=list)
     activations: list[CtrActivation] = field(default_factory=list)
     tags: list[str] = field(default_factory=list)
-
-
-@dataclass
-class HyperEdgeImport(HyperEdge):
-    name: str
-    imported_name: VarOrParam
-    imported_from: str
-    redefinitions: list[Definition]
\ No newline at end of file

src/gboml/ast/nodes.py

@@ -24,7 +24,7 @@ class NodeDefinition(Node):
     parameters: list[Definition] = field(default_factory=list)
     nodes: list[Node] = field(default_factory=list)
     hyperedges: list[HyperEdge] = field(default_factory=list)
-    variables: list[VariableDefinition] = field(default_factory=list)
+    variables: list[VariableDefinition | ScopeChange] = field(default_factory=list)
     constraints: list[Constraint] = field(default_factory=list)
     objectives: list[Objective] = field(default_factory=list)
     activations: list[Activation] = field(default_factory=list)
@@ -39,17 +39,8 @@ class NodeGenerator(Node):
     parameters: list[Definition] = field(default_factory=list)
     nodes: list[Node] = field(default_factory=list)
     hyperedges: list[HyperEdge] = field(default_factory=list)
-    variables: list[VariableDefinition] = field(default_factory=list)
+    variables: list[VariableDefinition | ScopeChange] = field(default_factory=list)
     constraints: list[Constraint] = field(default_factory=list)
     objectives: list[Objective] = field(default_factory=list)
     activations: list[Activation] = field(default_factory=list)
     tags: list[str] = field(default_factory=list)
-
-
-@dataclass
-class NodeImport(Node):
-    name: str
-    imported_name: VarOrParam
-    imported_from: str
-    scope_changes: list[ScopeChange] = field(default_factory=list)
-    redefinitions: list[Definition] = field(default_factory=list)

src/gboml/parsing.py

@@ -63,7 +63,6 @@ def _lark_to_gboml(tree: Tree, filename: Optional[str] = None) -> GBOMLGraph:
         # obj(*children, meta=meta)
         #
         to_obj = {
-            "hyperedge_import": HyperEdgeImport,
             "var_or_param_leaf": VarOrParamLeaf,
             "var_or_param": VarOrParam,
             "constraint_std": StdConstraint,
@@ -173,9 +172,14 @@ def _lark_to_gboml(tree: Tree, filename: Optional[str] = None) -> GBOMLGraph:
                                      objectives_block, activations, tags, meta=meta)
 
         def node_import(self, meta: Meta, name: str, imported_name: VarOrParam, imported_from: str, redef: list[ScopeChange | Definition]):
-            return NodeImport(name, imported_name, imported_from,
-                              [x for x in redef if isinstance(x, ScopeChange)],
-                              [x for x in redef if isinstance(x, Definition)], meta=meta)
+            return NodeDefinition(name, Extends(imported_name, imported_from, meta=meta),
+                                  parameters=[x for x in redef if isinstance(x, Definition)],
+                                  variables=[x for x in redef if isinstance(x, ScopeChange)],
+                                  meta=meta)
+
+        def hyperedge_import(self, meta: Meta, name: str, imported_name: VarOrParam, imported_from: str, redef: list[Definition]):
+            return HyperEdgeDefinition(name, Extends(imported_name, imported_from, meta=meta),
+                                       parameters=redef, meta=meta)
 
         def start(self, meta: Meta, time_horizon: Optional[int], global_defs: list[Definition], nodes_hyperedges: NodesAndHyperEdges):
             return GBOMLGraph(time_horizon, global_defs, nodes_hyperedges.nodes, nodes_hyperedges.hyperedges, meta=meta)