Enable Typing and generics based typing to support PEP585 (#727)

* Enable Typing and generics based types to support PEP585

Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>

* Additional fixes

Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>

* More fixes

Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>

---------

Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>

Author: arhamchopra

csp/impl/types/container_type_normalizer.py

@@ -1,3 +1,5 @@
+import collections.abc
+import types
 import typing
 
 import numpy
@@ -23,11 +25,78 @@ class ContainerTypeNormalizer:
         csp.typing.NumpyNDArray: numpy.ndarray,
     }
 
+    @classmethod
+    def canonicalize_builtin_generics(cls, typ):
+        """Recursively canonicalize PEP 585 builtin generics to their ``typing`` equivalents so that,
+        e.g., ``list[int]`` and ``typing.List[int]`` (and every nesting combination) normalize to a single
+        representation that compares and hashes equal.
+
+        Only the builtin container origins in ``_ORIGIN_COMPAT_MAP`` (list, set, dict, tuple) are remapped
+        to their ``typing`` form. Every other generic wrapper (FastList, csp.typing numpy arrays,
+        typing.Callable, typing.Mapping, custom generics, ...) keeps its own origin/flavor, but we still
+        recurse into its arguments so nested builtin containers get canonicalized. Unions (both
+        ``typing.Union``/``Optional`` and PEP 604 ``X | Y``) are traversed as well. When nothing actually
+        changes, the original object is returned unchanged, both to avoid needless allocations and to
+        preserve object identity for callers that rely on it.
+        """
+        if CspTypingUtils.is_union_type(typ):
+            args = typing.get_args(typ)
+            converted_args = tuple(cls._canonicalize_arg(arg) for arg in args)
+            if converted_args == args:
+                return typ
+            return typing.Union[converted_args]
+
+        if CspTypingUtils.is_generic_container(typ):
+            # __args__ (rather than get_args) keeps Callable's flattened ([params], ret) shape, which is
+            # what copy_with expects for faithful reconstruction.
+            args = typ.__args__
+            converted_args = tuple(cls._canonicalize_arg(arg) for arg in args)
+            canonical_origin = CspTypingUtils._ORIGIN_COMPAT_MAP.get(typ.__origin__)
+            if canonical_origin is not None:
+                # list/set/dict/tuple: rewrite to the typing form. A builtin (types.GenericAlias) alias
+                # must always be rebuilt (that is the whole point); an already-typing alias whose args did
+                # not change is returned as-is to preserve identity.
+                if converted_args == args and not isinstance(typ, types.GenericAlias):
+                    return typ
+                return canonical_origin[converted_args if len(converted_args) != 1 else converted_args[0]]
+            # Preserved wrapper (FastList, numpy arrays, Callable, Mapping, custom generic, ...): keep the
+            # outer origin/flavor but rebuild with canonicalized args when a child actually changed.
+            if converted_args == args:
+                return typ
+            if hasattr(typ, "copy_with"):
+                # typing._GenericAlias (typing.Callable, typing.Mapping, csp numpy arrays, ...): copy_with
+                # faithfully preserves the alias flavor, including Callable's flattened arg shape.
+                return typ.copy_with(converted_args)
+            origin = typ.__origin__
+            if origin is collections.abc.Callable:
+                # PEP 585 collections.abc.Callable[[p1, ...], ret]: __args__ is flattened, so restore the
+                # ([params], ret) subscription shape (an Ellipsis param list stays as Callable[..., ret]).
+                *params, ret = converted_args
+                if params == [Ellipsis]:
+                    return origin[..., ret]
+                return origin[list(params), ret]
+            try:
+                return origin[converted_args if len(converted_args) != 1 else converted_args[0]]
+            except TypeError:
+                # Exotic / non-subscriptable origin: leave it unchanged rather than fail normalization.
+                return typ
+
+        return typ
+
+    @classmethod
+    def _canonicalize_arg(cls, arg):
+        # A bare string argument inside a generic (e.g. ``list["T"]``) is stored raw by PEP 585 builtins but
+        # as a ForwardRef by typing generics; canonicalize to ForwardRef so both spellings match (and so we
+        # do not mint a fresh TypeVar on every call, which would defeat equality/caching).
+        if isinstance(arg, str):
+            return typing.ForwardRef(arg)
+        return cls.canonicalize_builtin_generics(arg)
+
     @classmethod
     def _convert_containers_to_typing_generic_meta(cls, typ, is_within_container):
+        typ = cls.canonicalize_builtin_generics(typ)
         if CspTypingUtils.is_generic_container(typ):
             return typ
-            # cls._deep_convert_generic_meta_to_typing_generic_meta(typ, is_within_container)
         elif isinstance(typ, dict):
             # warn(
             #     "Using {K: V} syntax for type declaration is deprecated. Use Dict[K, V] instead.",

csp/impl/types/instantiation_type_resolver.py

@@ -403,9 +403,9 @@ def _is_scalar_value_matching_spec(self, inp_def_type, arg):
                 if self._is_scalar_value_matching_spec(t, arg):
                     return True
         if isinstance(arg, SnapType):
-            return arg.ts_type.typ is inp_def_type
+            return arg.ts_type.typ == inp_def_type
         if isinstance(arg, SnapKeyType):
-            return arg.key_tstype.typ is inp_def_type
+            return arg.key_tstype.typ == inp_def_type
         return False
 
     def _rec_validate_container_and_resolve_tvars(self, sub_arg, sub_type_def):

csp/impl/types/pydantic_types.py

@@ -111,11 +111,11 @@ def make_snap_validator(inp_def_type):
 
     def snap_validator(v: Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo) -> Any:
         if isinstance(v, SnapType):
-            if v.ts_type.typ is inp_def_type:
+            if v.ts_type.typ == inp_def_type:
                 return v
             raise ValueError(f"Expecting {inp_def_type} for csp.snap value, but getting {v.ts_type.typ}")
         if isinstance(v, SnapKeyType):
-            if v.key_tstype.typ is inp_def_type:
+            if v.key_tstype.typ == inp_def_type:
                 return v
             raise ValueError(f"Expecting {inp_def_type} for csp.snap_key value, but getting {v.key_tstype.typ}")
         return handler(v)

csp/impl/wiring/base_parser.py

@@ -400,6 +400,12 @@ def parse_func_signature(self, funcdef):
             else:
                 typ = self._eval_expr(arg.annotation)
             arg_kind, basket_kind = self._resolve_input_type_kind(typ)
+            if arg_kind == ArgKind.SCALAR:
+                # Store the scalar note in canonical form once, at parse time, so downstream type
+                # comparisons (e.g. csp.snap) see list[X] and typing.List[X] as the same type without
+                # re-normalizing on every wiring. Deliberately narrower than normalize_type: numpy
+                # arrays, {K: V}/[T] shorthands, and bare aliases are intentionally left untouched.
+                typ = ContainerTypeNormalizer.canonicalize_builtin_generics(typ)
 
             inputs.append(InputDef(arg.arg, typ, arg_kind, basket_kind, tsidx, arg_idx))
 

csp/tests/impl/types/test_tstype.py

@@ -1,5 +1,21 @@
 import sys
-from typing import Any, Dict, ForwardRef, Generic, List, Mapping, TypeVar, Union, get_args, get_origin
+from datetime import datetime, timedelta
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    ForwardRef,
+    Generic,
+    List,
+    Mapping,
+    Optional,
+    Set,
+    Tuple,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+)
 from unittest import TestCase
 
 import numpy as np
@@ -9,9 +25,12 @@
 import csp
 from csp import dynamic_demultiplex, ts
 from csp.impl.types.common_definitions import OutputBasket, Outputs
+from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
 from csp.impl.types.pydantic_type_resolver import TVarValidationContext
 from csp.impl.types.pydantic_types import DynamicBasketPydantic
 from csp.impl.types.tstype import TsType
+from csp.impl.types.typing_utils import FastList
+from csp.typing import Numpy1DArray, NumpyNDArray
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -165,3 +184,117 @@ def test_validate(self):
         dynamic_basket = dynamic_demultiplex(csp.const(1.0), csp.const("A"))
         ta.validate_python(dynamic_basket)
         self.assertRaises(Exception, ta.validate_python, {csp.const("A"): csp.const(1.0)})
+
+
+class TestTsTypePep585Equivalence(TestCase):
+    """PEP 585 builtin generics (list[int], dict[str, int], ...) must normalize to the same TsType as
+    their typing equivalents (typing.List[int], typing.Dict[str, int], ...) so that tooling that rewrites
+    typing.List -> list (e.g. ruff UP006) does not change TsType equality/hashing."""
+
+    def test_equality_and_hash_parity(self):
+        cases = [
+            (list[int], List[int]),
+            (dict[str, int], Dict[str, int]),
+            (set[int], Set[int]),
+            (tuple[int, str], Tuple[int, str]),
+            (tuple[int, ...], Tuple[int, ...]),
+        ]
+        for builtin, typing_form in cases:
+            with self.subTest(builtin=builtin):
+                self.assertEqual(ts[builtin], ts[typing_form])
+                self.assertEqual(hash(ts[builtin]), hash(ts[typing_form]))
+
+    def test_nested_equality(self):
+        self.assertEqual(ts[list[dict[str, int]]], ts[List[Dict[str, int]]])
+        self.assertEqual(
+            ts[dict[str, list[tuple[int, set[str]]]]],
+            ts[Dict[str, List[Tuple[int, Set[str]]]]],
+        )
+        self.assertEqual(
+            hash(ts[list[dict[str, int]]]),
+            hash(ts[List[Dict[str, int]]]),
+        )
+
+    def test_canonicalizes_to_typing_form(self):
+        # The canonical inner type is the typing form (the representation csp inference already emits).
+        self.assertEqual(ts[list[int]].typ, List[int])
+        self.assertEqual(ts[dict[str, int]].typ, Dict[str, int])
+        self.assertEqual(ts[set[int]].typ, Set[int])
+        self.assertEqual(ts[tuple[int, str]].typ, Tuple[int, str])
+
+    def test_const_inference_matches_modernized_annotation(self):
+        # csp.const infers typing.List[...]; a modernized (ruff UP006) annotation is ts[list[...]].
+        # These previously compared unequal, which broke strict-equality checks (e.g. csp-gateway channels).
+        self.assertEqual(csp.const([1, 2, 3]).tstype, ts[list[int]])
+        self.assertEqual(csp.const({"a": 1}).tstype, ts[dict[str, int]])
+
+    def test_preserves_non_builtin_generics(self):
+        # FastList / csp numpy array types share machinery with typing generics but must NOT be collapsed.
+        self.assertIs(ContainerTypeNormalizer.normalize_type(FastList[int]).__origin__, FastList)
+        self.assertIs(ContainerTypeNormalizer.normalize_type(Numpy1DArray[float]).__origin__, Numpy1DArray)
+        self.assertIs(ContainerTypeNormalizer.normalize_type(NumpyNDArray[float]).__origin__, NumpyNDArray)
+
+    def test_node_binding_builtin_generic(self):
+        # A builtin-generic-annotated node input accepts an edge whose inferred type uses the typing form.
+        @csp.node
+        def consume(x: ts[list[int]]) -> ts[int]:
+            if csp.ticked(x):
+                return len(x)
+
+        @csp.graph
+        def g():
+            csp.add_graph_output("o", consume(csp.const([1, 2, 3])))
+
+        results = csp.run(g, starttime=datetime(2020, 1, 1), endtime=timedelta(days=1))
+        self.assertEqual(results["o"][0][1], 3)
+
+    def test_builtin_generic_nested_in_union(self):
+        # A builtin container nested inside a union (typing.Optional or PEP 604 X | None) is canonicalized.
+        self.assertEqual(ts[list[int] | None], ts[Optional[List[int]]])
+        self.assertEqual(ts[Optional[list[int]]], ts[Optional[List[int]]])
+        self.assertEqual(ts[dict[str, list[int]] | None], ts[Optional[Dict[str, List[int]]]])
+        self.assertEqual(hash(ts[list[int] | None]), hash(ts[Optional[List[int]]]))
+
+    def test_builtin_generic_nested_in_preserved_wrapper(self):
+        # The outer wrapper (Mapping / FastList / Callable) is preserved, but builtin containers nested
+        # inside it are still canonicalized so equality holds after a typing.List -> list rewrite.
+        self.assertEqual(ts[Mapping[str, list[int]]], ts[Mapping[str, List[int]]])
+        self.assertEqual(ts[FastList[list[int]]], ts[FastList[List[int]]])
+        self.assertEqual(ts[Callable[[list[int]], dict[str, int]]], ts[Callable[[List[int]], Dict[str, int]]])
+        # ... and the wrapper origin is not collapsed into typing.List/dict/etc.
+        self.assertIs(ContainerTypeNormalizer.normalize_type(FastList[list[int]]).__origin__, FastList)
+
+    def test_string_arg_canonicalizes_to_forward_ref(self):
+        # PEP 585 builtins store a bare string arg (list["T"]) while typing stores a ForwardRef; both must
+        # normalize to the same stable ForwardRef (not a freshly-minted TypeVar on each call).
+        self.assertEqual(ts[list["T"]], ts[list["T"]])
+        self.assertEqual(ts[list["T"]], ts[List["T"]])
+        self.assertEqual(ContainerTypeNormalizer.normalize_type(list["T"]).__args__[0], ForwardRef("T"))
+
+    def test_identity_preserved_when_unchanged(self):
+        # Nothing to canonicalize -> return the very same object, since some call sites (csp.snap
+        # validation) compare normalized types with `is`.
+        self.assertIs(ContainerTypeNormalizer.normalize_type(List[int]), List[int])
+        self.assertIs(ContainerTypeNormalizer.normalize_type(List["Foo"]), List["Foo"])
+        self.assertIs(ContainerTypeNormalizer.normalize_type(Optional[int]), Optional[int])
+        self.assertIs(ContainerTypeNormalizer.normalize_type(Dict[str, List[int]]), Dict[str, List[int]])
+
+    def test_empty_tuple(self):
+        self.assertEqual(ts[tuple[()]], ts[Tuple[()]])
+
+    def test_callable_inner_canonicalized(self):
+        import collections.abc as abc
+
+        normalize = ContainerTypeNormalizer.normalize_type
+        # typing.Callable: builtin containers nested in params/return canonicalize identically for both
+        # spellings, so the modernized and classic annotations stay equal.
+        self.assertEqual(
+            normalize(Callable[[list[int]], dict[str, int]]),
+            normalize(Callable[[List[int]], Dict[str, int]]),
+        )
+        # collections.abc.Callable (the ruff UP006 rewrite of typing.Callable) must reconstruct without
+        # error for finite, empty and ellipsis parameter lists, canonicalizing nested builtins while
+        # keeping its own origin.
+        self.assertEqual(normalize(abc.Callable[[list[int]], str]), abc.Callable[[List[int]], str])
+        self.assertEqual(normalize(abc.Callable[[], list[int]]), abc.Callable[[], List[int]])
+        self.assertEqual(normalize(abc.Callable[..., list[int]]), abc.Callable[..., List[int]])

csp/tests/test_dynamic.py

@@ -1,18 +1,21 @@
 import itertools
+import os
 import random
 import string
 import time
 import unittest
 from collections import defaultdict
 from datetime import datetime, timedelta
-from typing import Dict, List
+from typing import Dict, List, Optional, Set
 
 import numpy
 
 import csp
 from csp import ts
 from csp.utils.datetime import utc_now
 
+USE_PYDANTIC = os.environ.get("CSP_PYDANTIC", True)
+
 
 class DynData(csp.Struct):
     key: str
@@ -188,7 +191,73 @@ def g():
             self.assertEqual(len(res[f"{key}_tsadj"]), ts_ticks)
             self.assertTrue(all(x[1].val * 2 == y[1] for x, y in zip(res[f"{key}_ts"], res[f"{key}_tsadj"])))
 
-    def test_shared_input(self):
+    def test_snap_builtin_generic_scalar(self):
+        # csp.snap of a container-typed edge into a scalar arg written with the PEP 585 builtin form
+        # (list[str] rather than typing.List[str]) must still type-check.
+        @csp.graph
+        def dyn_graph(key: str, snapped: list[str]):
+            csp.add_graph_output(f"{key}_snapped", csp.const(snapped))
+
+        def g():
+            keys = csp.curve(List[str], [(timedelta(seconds=1), ["A", "B"])])
+            basket = gen_basket(keys, csp.null_ts(List[str]))
+            csp.dynamic(basket, dyn_graph, csp.snapkey(), csp.snap(keys))
+
+        res = csp.run(g, starttime=datetime(2021, 6, 22), endtime=timedelta(seconds=3))
+        self.assertIn("A", res["A_snapped"][0][1])
+        self.assertIn("B", res["B_snapped"][0][1])
+
+    def test_snap_builtin_generic_dict_scalar(self):
+        # csp.snap of a dict-typed edge into a scalar arg written in the PEP 585 builtin form
+        # (dict[str, int] rather than typing.Dict[str, int]) must still type-check.
+        @csp.graph
+        def dyn_graph(key: str, snapped: dict[str, int]):
+            csp.add_graph_output(f"{key}_snapped", csp.const(snapped))
+
+        def g():
+            keys = csp.curve(List[str], [(timedelta(seconds=1), ["A", "B"])])
+            values = csp.curve(Dict[str, int], [(timedelta(seconds=1), {"A": 1, "B": 2})])
+            basket = gen_basket(keys, csp.null_ts(List[str]))
+            csp.dynamic(basket, dyn_graph, csp.snapkey(), csp.snap(values))
+
+        res = csp.run(g, starttime=datetime(2021, 6, 22), endtime=timedelta(seconds=3))
+        self.assertEqual(res["A_snapped"][0][1], {"A": 1, "B": 2})
+        self.assertEqual(res["B_snapped"][0][1], {"A": 1, "B": 2})
+
+    def test_snap_builtin_generic_set_scalar(self):
+        # csp.snap of a set-typed edge into a scalar arg written in the PEP 585 builtin form
+        # (set[str] rather than typing.Set[str]) must still type-check.
+        @csp.graph
+        def dyn_graph(key: str, snapped: set[str]):
+            csp.add_graph_output(f"{key}_snapped", csp.const(snapped))
+
+        def g():
+            keys = csp.curve(List[str], [(timedelta(seconds=1), ["A", "B"])])
+            values = csp.curve(Set[str], [(timedelta(seconds=1), {"X", "Y"})])
+            basket = gen_basket(keys, csp.null_ts(List[str]))
+            csp.dynamic(basket, dyn_graph, csp.snapkey(), csp.snap(values))
+
+        res = csp.run(g, starttime=datetime(2021, 6, 22), endtime=timedelta(seconds=3))
+        self.assertEqual(res["A_snapped"][0][1], {"X", "Y"})
+        self.assertEqual(res["B_snapped"][0][1], {"X", "Y"})
+
+    @unittest.skipIf(USE_PYDANTIC, "csp.snap into a union-typed scalar is only supported by the legacy type resolver")
+    def test_snap_builtin_generic_union_scalar(self):
+        # A builtin generic nested inside a union scalar annotation (Optional[list[str]]) exercises
+        # the union branch of the legacy type resolver, which must normalize list[str] to typing.List[str]
+        # for the snapped edge type to match.
+        @csp.graph
+        def dyn_graph(key: str, snapped: Optional[list[str]]):
+            csp.add_graph_output(f"{key}_snapped", csp.const(snapped))
+
+        def g():
+            keys = csp.curve(List[str], [(timedelta(seconds=1), ["A", "B"])])
+            basket = gen_basket(keys, csp.null_ts(List[str]))
+            csp.dynamic(basket, dyn_graph, csp.snapkey(), csp.snap(keys))
+
+        res = csp.run(g, starttime=datetime(2021, 6, 22), endtime=timedelta(seconds=3))
+        self.assertIn("A", res["A_snapped"][0][1])
+        self.assertIn("B", res["B_snapped"][0][1])
         """ensure an externally wired input is shared / not recreated per sub-graph"""
         instances = []