parser fix

Signed-off-by: Mark Halka <Mark.Halka@Point72.com>

Author: markhalka

csp/impl/wiring/csp_numba/csp_node_transformer.py

@@ -1,8 +1,9 @@
 import ast
-import inspect
-import textwrap
 from dataclasses import dataclass, field
-from typing import List, Optional, Set, Union
+from typing import TYPE_CHECKING, List, Optional, Union
+
+if TYPE_CHECKING:
+    from csp.impl.wiring.node_parser import ParsedNodeDefinition
 
 
 @dataclass
@@ -23,6 +24,7 @@ class TransformedNode:
     transformed_source: str = ""
     start_body: List[ast.AST] = field(default_factory=list)  # Code from with csp.start():
     stop_body: List[ast.AST] = field(default_factory=list)  # Code from with csp.stop():
+    transformed_ast: Optional[ast.FunctionDef] = None
 
 
 class CspNodeTransformer(ast.NodeTransformer):
@@ -41,7 +43,6 @@ class CspNodeTransformer(ast.NodeTransformer):
 
     def __init__(self):
         self.state_variables: List[StateVariable] = []
-        self.input_names: Set[str] = set()
         self.start_body: List[ast.AST] = []
         self.stop_body: List[ast.AST] = []
         self.csp_call_transformers = {
@@ -56,38 +57,6 @@ def __init__(self):
                 op=ast.And(),
             ),
         }
-        self.special_block_transformers = {
-            "state": self._transform_state_block,
-            "start": lambda node: self.start_body.extend(self._flatten_transformed_statements(node.body)) or [],
-            "stop": lambda node: self.stop_body.extend(self._flatten_transformed_statements(node.body)) or [],
-            "alarms": lambda node: [],
-        }
-
-    def _is_special_block(self, node: ast.AST, block_name: str) -> bool:
-        """Returns True if this is a block we need to handle"""
-        if not isinstance(node, ast.With):
-            return False
-
-        if len(node.items) != 1:
-            return False
-
-        context_expr = node.items[0].context_expr
-
-        if isinstance(context_expr, ast.Call):
-            func = context_expr.func
-            if isinstance(func, ast.Name) and func.id == block_name:
-                return True
-            if isinstance(func, ast.Attribute):
-                if isinstance(func.value, ast.Name) and func.value.id == "csp" and func.attr == block_name:
-                    return True
-
-        return False
-
-    def _get_special_block_name(self, node: ast.AST) -> Optional[str]:
-        for block_name in self.special_block_transformers:
-            if self._is_special_block(node, block_name):
-                return block_name
-        return None
 
     def _infer_type(self, name: str, value: ast.AST) -> ast.AST:
         if isinstance(value, ast.Constant):
@@ -203,21 +172,10 @@ def _create_state_assignment(
         self.state_variables.append(StateVariable(name=name, type_annotation=type_annotation, initial_value=value))
         return ast.AnnAssign(target=target, annotation=state_type, value=value, simple=simple)
 
-    def _transform_state_block(self, node: ast.With) -> List[ast.AST]:
-        """
-        Transform state block variables to have State[type] annotations.
-
-        with csp.state():
-            x = 0
-            y: float = 1.0
-
-        becomes:
-            x: State[int] = 0
-            y: State[float] = 1.0
-        """
+    def _transform_state_statements(self, statements) -> List[ast.AST]:
         transformed = []
 
-        for stmt in node.body:
+        for stmt in statements:
             if isinstance(stmt, ast.AnnAssign):
                 # Already annotated: y: float = 1.0
                 name = stmt.target.id if isinstance(stmt.target, ast.Name) else None
@@ -263,31 +221,12 @@ def _flatten_transformed_statements(self, statements: List[ast.AST]) -> List[ast
                 flattened.append(transformed_stmt)
         return flattened
 
-    def _transform_body(self, body: List[ast.AST]) -> List[ast.AST]:
-        transformed = []
-
-        for stmt in body:
-            block_name = self._get_special_block_name(stmt)
-            if block_name is not None:
-                transformed.extend(self.special_block_transformers[block_name](stmt))
-                continue
-
-            # Transform the statement
-            transformed.extend(self._flatten_transformed_statements([stmt]))
-
-        return transformed
-
-    def _transform_func_def(self, func_def: ast.FunctionDef) -> TransformedNode:
+    def _reset(self):
         self.state_variables = []
-        self.input_names = set()
         self.start_body = []
         self.stop_body = []
 
-        # Collect input names (don't transform annotations)
-        for arg in func_def.args.args:
-            self.input_names.add(arg.arg)
-
-        transformed_body = self._transform_body(func_def.body)
+    def _build_transformed_node(self, func_def: ast.FunctionDef, transformed_body: List[ast.AST]) -> TransformedNode:
         new_func_def = ast.FunctionDef(
             name=func_def.name,
             args=func_def.args,  # Keep original args with ts[type] annotations
@@ -304,21 +243,16 @@ def _transform_func_def(self, func_def: ast.FunctionDef) -> TransformedNode:
             state_variables=self.state_variables,
             transformed_body=transformed_body,
             original_ast=func_def,
+            transformed_ast=new_func_def,
             transformed_source=transformed_source,
             start_body=self.start_body,
             stop_body=self.stop_body,
         )
 
-    def transform_csp_node(self, func) -> TransformedNode:
-        if callable(func):
-            source = textwrap.dedent(inspect.getsource(func))
-        else:
-            source = textwrap.dedent(func)
-
-        tree = ast.parse(source)
-        func_def = tree.body[0]
-
-        if not isinstance(func_def, ast.FunctionDef):
-            raise ValueError("Expected a function definition")
-
-        return self._transform_func_def(func_def)
+    def transform_csp_node(self, definition: "ParsedNodeDefinition") -> TransformedNode:
+        self._reset()
+        self.start_body = self._flatten_transformed_statements(definition.blocks.start)
+        self.stop_body = self._flatten_transformed_statements(definition.blocks.stop)
+        transformed_body = self._transform_state_statements(definition.blocks.state)
+        transformed_body.extend(self._flatten_transformed_statements(definition.blocks.body))
+        return self._build_transformed_node(definition.funcdef, transformed_body)

csp/impl/wiring/csp_numba/input_handlers.py

@@ -1,165 +1,29 @@
-"""
-CSP timeseries input handlers for numba_node.
-
-Provides:
-- TsInputHandler for ts[type] parameters
-- ListBasketInputHandler for [ts[type]] list basket parameters
-- DictBasketInputHandler for {key: ts[type]} dict basket parameters
-"""
+"""Pass CSP-parsed timeseries input metadata to numba_cfunc_compiler."""
 
 import inspect
 from dataclasses import dataclass
-from typing import Any, Optional, get_args, get_origin
+from typing import Any, Optional
 
 from numba_cfunc_compiler.function_analyzer import InputTypeHandler
 from numba_cfunc_compiler.models import ParameterInfo
 
-from csp.impl.types.tstype import TsType
-from csp.impl.wiring.edge import Edge
-
-
-def _extract_ts_inner_type(ann: Any) -> Optional[type]:
-    if get_origin(ann) is not TsType:
-        return None
-    args = get_args(ann)
-    return args[0] if args else getattr(ann, "typ", None)
-
-
-def _parse_list_basket_annotation(ann: Any) -> tuple[bool, Optional[type]]:
-    # Legacy syntax: [ts[T]]
-    if isinstance(ann, list):
-        if len(ann) != 1:
-            return True, None
-        return True, _extract_ts_inner_type(ann[0])
-
-    # Modern syntax: list[ts[T]] / typing.List[ts[T]]
-    if get_origin(ann) is list:
-        args = get_args(ann)
-        if len(args) != 1:
-            return True, None
-        return True, _extract_ts_inner_type(args[0])
-
-    return False, None
-
 
 @dataclass(frozen=True)
-class _DictBasketExpectedType:
-    key_type: type
-    element_type: type
-
+class CspInputMetadata:
+    category: str
 
-def _parse_dict_basket_annotation(ann: Any) -> tuple[bool, Optional[type], Optional[type]]:
-    key_type = None
-    value_ann = None
-
-    # Legacy syntax: {str: ts[T]} / {int: ts[T]}
-    if isinstance(ann, dict):
-        if len(ann) != 1:
-            return True, None, None
-        key_type, value_ann = next(iter(ann.items()))
-    # Modern syntax: dict[str, ts[T]] / typing.Dict[str, ts[T]]
-    elif get_origin(ann) is dict:
-        args = get_args(ann)
-        if len(args) != 2:
-            return True, None, None
-        key_type, value_ann = args
-    else:
-        return False, None, None
-
-    if key_type not in (str, int):
-        return False, None, None
-
-    return True, key_type, _extract_ts_inner_type(value_ann)
-
-
-def _validate_edge_type(param_name: str, edge: Edge, expected_type: type) -> Edge:
-    actual_type = edge.tstype.typ
-    if actual_type is not expected_type and not (expected_type is float and actual_type is int):
-        raise TypeError(
-            f"Argument '{param_name}' expected ts[{expected_type.__name__}], got ts[{actual_type.__name__}]"
-        )
-    return edge
-
-
-class TsInputHandler(InputTypeHandler):
-    """Handles ts[type] input annotations."""
 
+class CspInputHandler(InputTypeHandler):
     def try_parse(self, param: inspect.Parameter, ann: Any) -> Optional[ParameterInfo]:
-        inner_type = _extract_ts_inner_type(ann)
-        if inner_type is None:
+        if not isinstance(ann, CspInputMetadata):
             return None
-
-        return ParameterInfo(expected_type=inner_type, category="signal")
+        return ParameterInfo(expected_type=ann, category=ann.category)
 
     def validate_value(self, param_name: str, value: Any, expected_type: Any) -> Any:
-        if not isinstance(value, Edge):
-            raise TypeError(f"Argument '{param_name}' must be an Edge, got {type(value).__name__}")
-        return _validate_edge_type(param_name, value, expected_type)
-
-
-class ListBasketInputHandler(InputTypeHandler):
-    """Handles [ts[type]] and list[ts[type]] basket input annotations."""
-
-    def try_parse(self, param: inspect.Parameter, ann: Any) -> Optional[ParameterInfo]:
-        matched, inner_type = _parse_list_basket_annotation(ann)
-        if not matched:
-            return None
-        if inner_type is None:
-            raise TypeError(f"List basket '{param.name}' element ts[type] is missing type argument")
-
-        return ParameterInfo(expected_type=inner_type, category="signal_set")
-
-    def validate_value(self, param_name: str, value: Any, expected_type: Any) -> Any:
-        if not isinstance(value, (list, tuple)):
-            raise TypeError(f"Argument '{param_name}' must be a list, got {type(value).__name__}")
-        if not value:
-            raise ValueError(f"List basket '{param_name}' must not be empty")
-
-        result = {}
-        for i, edge in enumerate(value):
-            if not isinstance(edge, Edge):
-                raise TypeError(f"List basket '{param_name}[{i}]' must be an Edge, got {type(edge).__name__}")
-            result[i] = _validate_edge_type(f"{param_name}[{i}]", edge, expected_type)
-        return result
-
-
-class DictBasketInputHandler(InputTypeHandler):
-    """Handles {key_type: ts[type]} and dict[key_type, ts[type]] basket annotations."""
-
-    def try_parse(self, param: inspect.Parameter, ann: Any) -> Optional[ParameterInfo]:
-        matched, key_type, inner_type = _parse_dict_basket_annotation(ann)
-        if not matched:
-            return None
-        if inner_type is None:
-            raise TypeError(f"Dict basket '{param.name}' element ts[type] is missing type argument")
-
-        return ParameterInfo(
-            expected_type=_DictBasketExpectedType(key_type=key_type, element_type=inner_type),
-            category="signal_set",
-        )
-
-    def validate_value(self, param_name: str, value: Any, expected_type: Any) -> Any:
-        if not isinstance(value, dict):
-            raise TypeError(f"Argument '{param_name}' must be a dict, got {type(value).__name__}")
-        if not value:
-            raise ValueError(f"Dict basket '{param_name}' must not be empty")
-
-        result = {}
-        for key, edge in value.items():
-            if not isinstance(key, expected_type.key_type):
-                raise TypeError(
-                    f"Dict basket '{param_name}' key {key!r} must be "
-                    f"{expected_type.key_type.__name__}, got {type(key).__name__}"
-                )
-            if not isinstance(edge, Edge):
-                raise TypeError(f"Dict basket '{param_name}[{key!r}]' must be an Edge, got {type(edge).__name__}")
-            result[key] = _validate_edge_type(f"{param_name}[{key!r}]", edge, expected_type.element_type)
-        return result
+        return value
 
 
 def register():
     from numba_cfunc_compiler.function_analyzer import FunctionAnalyzer
 
-    FunctionAnalyzer.register_input_handler(TsInputHandler())
-    FunctionAnalyzer.register_input_handler(ListBasketInputHandler())
-    FunctionAnalyzer.register_input_handler(DictBasketInputHandler())
+    FunctionAnalyzer.register_input_handler(CspInputHandler())