cleanup

Author: ducky64

edg/core/Builder.py

@@ -14,12 +14,17 @@ class Builder:
   def __init__(self) -> None:
     self.stack: List[BaseBlock] = []
 
-  def push_element(self, elt: BaseBlock) -> None:
-    self.stack.append(elt)
+  def push_element(self, elt: BaseBlock) -> Optional[BaseBlock]:
+    """Pushes a new element onto the context stack, returning the previous top element.
+    Ignores if the element is already on top of the stack."""
+    prev_elt = self.get_enclosing_block()
+    if not self.stack or self.stack[-1] is not elt:  # prevent double-pushing
+      self.stack.append(elt)
+    return prev_elt
 
   def pop_to(self, elt: Optional[BaseBlock]) -> None:
-    self.stack.pop()
-    assert self.get_enclosing_block() is elt
+    while (elt is None and self.stack) or (elt is not None and self.stack[-1] is not elt):
+      self.stack.pop()
 
   def get_enclosing_block(self) -> Optional[BaseBlock]:
     if not self.stack:

edg/core/HierarchyBlock.py

@@ -181,14 +181,7 @@ def remap_arg(arg_name: str, arg_type: Type[ConstraintExpr], arg_value: Any) ->
 
           return arg_type()._bind(InitParamBinding(self, typed_arg_value))
 
-        if builder.get_enclosing_block() is not self:
-          # test needed to make sure we don't double-push in nested super().__init__ calls
-          # create wrapper ConstraintExpr in new object scope
-          builder_prev = (True, builder.get_enclosing_block())
-          builder.push_element(self)
-        else:
-          builder_prev = (False, None)  # dummy value for try/finally
-
+        builder_prev = builder.push_element(self)
         try:
           # rebuild args and kwargs by traversing the args list
           new_args: List[Any] = []
@@ -228,8 +221,7 @@ def remap_arg(arg_name: str, arg_type: Type[ConstraintExpr], arg_value: Any) ->
 
           orig_init(self, *new_args, **new_kwargs)
         finally:
-          if builder_prev[0]:
-            builder.pop_to(builder_prev[1])
+          builder.pop_to(builder_prev)
 
       new_cls.__init__ = functools.update_wrapper(wrapped_init, orig_init)
 

edg/core/Link.py

@@ -20,19 +20,11 @@ def __new__(cls, *args: Any, **kwargs: Any) -> Any:
     if '__init__' in new_cls.__dict__:
       orig_init = new_cls.__dict__['__init__']
       def wrapped_init(self, *args, **kwargs) -> None:
-        if builder.get_enclosing_block() is not self:
-          # test needed to make sure we don't double-push in nested super().__init__ calls
-          # create wrapper ConstraintExpr in new object scope
-          builder_prev = (True, builder.get_enclosing_block())
-          builder.push_element(self)
-        else:
-          builder_prev = (False, None)  # dummy value for try/finally
-
+        builder_prev = builder.push_element(self)
         try:
           orig_init(self, *args, **kwargs)
         finally:
-          if builder_prev[0]:
-            builder.pop_to(builder_prev[1])
+          builder.pop_to(builder_prev)
 
       new_cls.__init__ = functools.update_wrapper(wrapped_init, orig_init)
 

edg/core/Ports.py

@@ -14,6 +14,7 @@
 
 if TYPE_CHECKING:
   from .Blocks import BaseBlock
+  from .Link import Link
   from .PortBlocks import PortBridge, PortAdapter