Update for autocomplete.

Signed-off-by: James Goppert <james.goppert@gmail.com>

Author: jgoppert

cyecca/dsl/__init__.py

@@ -170,7 +170,7 @@
 from cyecca.dsl.equations import Assignment, Equation, IfEquation, IfEquationBranch, Reinit, WhenClause
 
 # Expression tree
-from cyecca.dsl.expr import Expr, ExprKind
+from cyecca.dsl.expr import Expr, ExprKind, ExprLike
 
 # Flat model representation
 from cyecca.dsl.flat_model import FlatModel
@@ -221,7 +221,7 @@
     terminal,
 )
 from cyecca.dsl.simulation import SimulationResult, Simulator
-from cyecca.dsl.types import DType, Indices, Shape, Var, VarKind
+from cyecca.dsl.types import DType, Indices, NumericValue, Shape, Var, VarKind
 
 # Variables
 from cyecca.dsl.variables import ArrayDerivativeExpr, DerivativeExpr, SymbolicVar, TimeVar
@@ -246,6 +246,7 @@
     "DType",
     "Shape",
     "Indices",
+    "NumericValue",
     "submodel",
     # Free functions (continuous)
     "der",
@@ -291,6 +292,7 @@
     "FlatModel",
     "Expr",
     "ExprKind",
+    "ExprLike",
     "Equation",
     # Causality analysis
     "analyze_causality",

cyecca/dsl/context.py

@@ -31,12 +31,13 @@
 from __future__ import annotations
 
 import threading
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
+from typing import TYPE_CHECKING, Callable, List, Optional, Union
 
 from beartype import beartype
 
 from cyecca.dsl.equations import IfEquation, IfEquationBranch, Reinit, WhenClause
-from cyecca.dsl.expr import Expr, to_expr
+from cyecca.dsl.expr import Expr, ExprLike, to_expr
+from cyecca.dsl.instance import SubmodelProxy
 from cyecca.dsl.variables import SymbolicVar
 
 if TYPE_CHECKING:
@@ -427,7 +428,7 @@ def to_when_clause(self) -> "WhenClause":
 
 
 @beartype
-def when(condition: Any) -> WhenContext:
+def when(condition: ExprLike) -> WhenContext:
     """
     Create a when-clause for event handling (Modelica MLS 8.5).
 
@@ -453,7 +454,7 @@ def _(m):
 
     Parameters
     ----------
-    condition : Expr or SymbolicVar
+    condition : ExprLike
         Boolean condition that triggers the when-clause
 
     Returns
@@ -474,7 +475,7 @@ def _(m):
 
 
 @beartype
-def reinit(var: SymbolicVar, expr: Any) -> Optional[Reinit]:
+def reinit(var: SymbolicVar, expr: ExprLike) -> Optional[Reinit]:
     """
     Reinitialize a continuous-time state variable at an event.
 
@@ -489,7 +490,7 @@ def reinit(var: SymbolicVar, expr: Any) -> Optional[Reinit]:
     ----------
     var : SymbolicVar
         The state variable to reinitialize
-    expr : Expr or numeric
+    expr : ExprLike
         The new value expression (can use pre() for previous values)
 
     Returns
@@ -529,7 +530,7 @@ def reinit(var: SymbolicVar, expr: Any) -> Optional[Reinit]:
 
 
 @beartype
-def connect(a: Any, b: Any) -> None:
+def connect(a: SubmodelProxy, b: SubmodelProxy) -> None:
     """
     Connect two connectors, generating appropriate connection equations.
 
@@ -598,19 +599,12 @@ def _(m):
     where positive flow is into the component.
     """
     from cyecca.dsl.equations import Equation
-    from cyecca.dsl.instance import SubmodelProxy
 
     ctx = get_current_equation_context()
     if ctx is None:
         raise RuntimeError("connect() can only be used inside an @equations block")
 
-    # Validate that both are SubmodelProxy (connector instances)
-    if not isinstance(a, SubmodelProxy):
-        raise TypeError(f"connect() first argument must be a connector, got {type(a)}")
-    if not isinstance(b, SubmodelProxy):
-        raise TypeError(f"connect() second argument must be a connector, got {type(b)}")
-
-    # Get the connector metadata
+    # Get the connector metadata (beartype already validates SubmodelProxy type)
     a_instance = a._instance
     b_instance = b._instance
     a_metadata = a_instance._metadata
@@ -828,7 +822,7 @@ def execute_equations_method(
 
 
 @beartype
-def if_eq(condition: Any) -> IfContext:
+def if_eq(condition: ExprLike) -> IfContext:
     """
     Create an if-equation for conditional equations (Modelica MLS 8.3.4).
 
@@ -855,7 +849,7 @@ def _(m):
 
     Parameters
     ----------
-    condition : Expr or SymbolicVar
+    condition : ExprLike
         Boolean condition that selects which equations are active
 
     Returns
@@ -877,7 +871,7 @@ def _(m):
 
 
 @beartype
-def elseif_eq(condition: Any) -> IfContext:
+def elseif_eq(condition: ExprLike) -> IfContext:
     """
     Create an elseif branch for an if-equation.
 
@@ -894,7 +888,7 @@ def _(m):
 
     Parameters
     ----------
-    condition : Expr or SymbolicVar
+    condition : ExprLike
         Boolean condition for this elseif branch
 
     Returns

cyecca/dsl/decorators.py

@@ -32,11 +32,29 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Type, TypeVar, Union, cast
+
+# Python 3.11+ has dataclass_transform in typing, earlier versions need typing_extensions
+try:
+    from typing import dataclass_transform
+except ImportError:
+    try:
+        from typing_extensions import dataclass_transform
+    except ImportError:
+        # Fallback: no-op decorator if not available
+        def dataclass_transform(**kwargs):  # type: ignore[misc]
+            def decorator(cls_or_fn):  # type: ignore[no-untyped-def]
+                return cls_or_fn
+
+            return decorator
+
 
 import numpy as np
 from beartype import beartype
 
+# TypeVar for preserving class type through decorators
+_T = TypeVar("_T")
+
 from cyecca.dsl.context import (
     execute_algorithm_method,
     execute_equations_method,
@@ -47,10 +65,15 @@
 )
 from cyecca.dsl.equations import ArrayEquation, Assignment, Equation, WhenClause
 from cyecca.dsl.instance import ModelInstance
-from cyecca.dsl.types import DType, Shape, SubmodelField, Var, VarKind
+from cyecca.dsl.types import DType, NumericValue, Shape, SubmodelField, Var, VarKind
+from cyecca.dsl.variables import SymbolicVar
 
+# For IDE autocomplete: Real/Integer/Boolean/String return SymbolicVar for type checking
+# but Var at runtime. This allows m.theta to autocomplete as SymbolicVar.
 if TYPE_CHECKING:
-    pass
+    _VarReturn = SymbolicVar
+else:
+    _VarReturn = Var
 
 
 # =============================================================================
@@ -114,7 +137,7 @@ def var(
     protected: bool = False,
     # Connector prefixes (Modelica MLS Ch. 9)
     flow: bool = False,
-) -> Var:
+) -> _VarReturn:
     """
     Declare a variable in a Cyecca model.
 
@@ -231,7 +254,7 @@ def Real(
     protected: bool = False,
     # Connector
     flow: bool = False,
-) -> Var:
+) -> _VarReturn:
     """
     Declare a Real (floating-point) variable.
 
@@ -326,7 +349,7 @@ def Integer(
     constant: bool = False,
     # Visibility
     protected: bool = False,
-) -> Var:
+) -> _VarReturn:
     """
     Declare an Integer variable.
 
@@ -408,7 +431,7 @@ def Boolean(
     constant: bool = False,
     # Visibility
     protected: bool = False,
-) -> Var:
+) -> _VarReturn:
     """
     Declare a Boolean variable.
 
@@ -480,7 +503,7 @@ def String(
     constant: bool = False,
     # Visibility
     protected: bool = False,
-) -> Var:
+) -> _VarReturn:
     """
     Declare a String variable.
 
@@ -536,7 +559,7 @@ def String(
 
 
 @beartype
-def submodel(model_class: Type, **overrides: Any) -> SubmodelField:
+def submodel(model_class: Type, **overrides: NumericValue) -> SubmodelField:
     """
     Declare a submodel (nested model) with optional parameter overrides.
 
@@ -547,7 +570,7 @@ def submodel(model_class: Type, **overrides: Any) -> SubmodelField:
     ----------
     model_class : Type
         The model class to instantiate as a submodel
-    **overrides : Any
+    **overrides : NumericValue
         Parameter value overrides. The parameter names must match
         parameters defined in the submodel class.
 
@@ -575,8 +598,9 @@ def submodel(model_class: Type, **overrides: Any) -> SubmodelField:
 # =============================================================================
 
 
+@dataclass_transform(field_specifiers=(Real, Integer, Boolean, String, var))
 @beartype
-def model(cls: Type[Any]) -> Type[Any]:
+def model(cls: Type[_T]) -> Type[_T]:
     """
     Decorator to convert a class into a Cyecca model.
 
@@ -670,6 +694,8 @@ class ModelClass(ModelInstance):
         __name__ = cls.__name__
         __qualname__ = cls.__qualname__
         __module__ = cls.__module__
+        # Copy annotations from original class for IDE autocomplete
+        __annotations__ = getattr(cls, "__annotations__", {})
 
         _equations_methods = equations_methods
         _initial_equations_methods = initial_equations_methods
@@ -704,16 +730,17 @@ def get_algorithm(self) -> List[Assignment]:
 
     ModelClass._dsl_metadata = metadata
 
-    return ModelClass
+    return cast(Type[_T], ModelClass)
 
 
 # =============================================================================
 # @function decorator - Modelica functions (Ch. 12)
 # =============================================================================
 
 
+@dataclass_transform(field_specifiers=(Real, Integer, Boolean, String, var))
 @beartype
-def function(cls: Type[Any]) -> Type[Any]:
+def function(cls: Type[_T]) -> Type[_T]:
     """
     Decorator to convert a class into a Cyecca function.
 
@@ -787,8 +814,9 @@ def get_function_metadata(self) -> FunctionMetadata:
 # =============================================================================
 
 
+@dataclass_transform(field_specifiers=(Real, Integer, Boolean, String, var))
 @beartype
-def block(cls: Type[Any]) -> Type[Any]:
+def block(cls: Type[_T]) -> Type[_T]:
     """
     Decorator to convert a class into a Cyecca block.
 
@@ -832,8 +860,9 @@ def block(cls: Type[Any]) -> Type[Any]:
 # =============================================================================
 
 
+@dataclass_transform(field_specifiers=(Real, Integer, Boolean, String, var))
 @beartype
-def connector(cls: Type[Any]) -> Type[Any]:
+def connector(cls: Type[_T]) -> Type[_T]:
     """
     Decorator to convert a class into a Cyecca connector.
 

cyecca/dsl/expr.py

@@ -31,16 +31,37 @@
 
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, Generator, Optional, Tuple
+from typing import TYPE_CHECKING, Generator, List, Optional, Tuple, Union
 
 import numpy as np
 from beartype import beartype
 
 from cyecca.dsl.types import Indices, Shape
 
 if TYPE_CHECKING:
+    from cyecca.dsl.algorithm import AlgorithmVar
     from cyecca.dsl.variables import DerivativeExpr, SymbolicVar, TimeVar
 
+# Type alias for anything that can be converted to an Expr
+# Used in operators and functions that accept expressions
+# Note: At runtime, beartype uses object (accepting anything) and to_expr() validates.
+# For static type checking, the Union provides proper type hints.
+if TYPE_CHECKING:
+    ExprLike = Union[
+        "Expr",
+        "SymbolicVar",
+        "DerivativeExpr",
+        "TimeVar",
+        "AlgorithmVar",
+        float,
+        int,
+        List["ExprLike"],
+        np.ndarray,
+    ]
+else:
+    # At runtime, accept any object - to_expr() will validate and convert
+    ExprLike = object
+
 
 class ExprKind(Enum):
     """Kinds of expression nodes."""
@@ -330,7 +351,7 @@ def __hash__(self) -> int:
 
 
 @beartype
-def to_expr(x: Any) -> Expr:
+def to_expr(x: ExprLike) -> Expr:
     """Convert various types to Expr."""
     if isinstance(x, Expr):
         return x