Fix to allow Literal and Union via | in csp annotations

Signed-off-by: Nijat Khanbabayev <[email protected]>

Author: NeejWeej

csp/impl/types/pydantic_types.py

@@ -1,7 +1,7 @@
 import sys
 import types
 import typing
-from typing import Any, ForwardRef, Generic, Optional, Type, TypeVar, Union, get_args, get_origin
+from typing import Any, ForwardRef, Generic, Literal, Optional, Type, TypeVar, Union, get_args, get_origin
 
 from pydantic import GetCoreSchemaHandler, ValidationInfo, ValidatorFunctionWrapHandler
 from pydantic_core import CoreSchema, core_schema
@@ -184,6 +184,8 @@ def adjust_annotations(
             return TsType[
                 adjust_annotations(args[0], top_level=False, in_ts=True, make_optional=False, forced_tvars=forced_tvars)
             ]
+        if origin is Literal:  # for literals, we stop converting
+            return Optional[annotation] if make_optional else annotation
         else:
             try:
                 if origin is CspTypeVar or origin is CspTypeVarType:

csp/impl/types/type_annotation_normalizer_transformer.py

@@ -51,6 +51,8 @@ def visit_arg(self, node):
         return node
 
     def visit_Subscript(self, node):
+        # We choose to avoid parsing here
+        # to maintain current behavior of allowing empty lists in our types
         return node
 
     def visit_List(self, node):
@@ -98,17 +100,13 @@ def visit_Call(self, node):
         return node
 
     def visit_Constant(self, node):
-        if not self._cur_arg:
-            return node
-
-        if self._cur_arg:
-            return ast.Call(
-                func=ast.Attribute(value=ast.Name(id="typing", ctx=ast.Load()), attr="TypeVar", ctx=ast.Load()),
-                args=[node],
-                keywords=[],
-            )
-        else:
+        if not self._cur_arg or not isinstance(node.value, str):
             return node
+        return ast.Call(
+            func=ast.Attribute(value=ast.Name(id="typing", ctx=ast.Load()), attr="TypeVar", ctx=ast.Load()),
+            args=[node],
+            keywords=[],
+        )
 
     def visit_Str(self, node):
         return self.visit_Constant(node)

csp/tests/impl/types/test_pydantic_types.py

@@ -1,6 +1,6 @@
 import sys
 from inspect import isclass
-from typing import Any, Callable, Dict, Generic, List, Optional, Type, TypeVar, Union, get_args, get_origin
+from typing import Any, Callable, Dict, Generic, List, Literal, Optional, Type, TypeVar, Union, get_args, get_origin
 from unittest import TestCase
 
 import csp
@@ -160,3 +160,12 @@ def test_force_tvars(self):
         self.assertAnnotationsEqual(
             adjust_annotations(CspTypeVarType[T], forced_tvars={"T": float}), Union[Type[float], Type[int]]
         )
+
+    def test_literal(self):
+        self.assertAnnotationsEqual(adjust_annotations(Literal["a", "b"]), Literal["a", "b"])
+        self.assertAnnotationsEqual(
+            adjust_annotations(Literal["a", "b"], make_optional=True), Optional[Literal["a", "b"]]
+        )
+        self.assertAnnotationsEqual(adjust_annotations(Literal[123, "a"]), Literal[123, "a"])
+        self.assertAnnotationsEqual(adjust_annotations(Literal[123, None]), Literal[123, None])
+        self.assertAnnotationsEqual(adjust_annotations(ts[Literal[123, None]]), ts[Literal[123, None]])

csp/tests/test_type_checking.py

@@ -1,6 +1,7 @@
 import os
 import pickle
 import re
+import sys
 import typing
 import unittest
 from datetime import datetime, time, timedelta
@@ -938,6 +939,60 @@ def test_is_callable(self):
             result = CspTypingUtils.is_callable(input_type)
             self.assertEqual(result, expected)
 
+    def test_literal_with_pipe_operator(self):
+        """Test combining Literal types with the pipe operator in Python 3.10+."""
+        if sys.version_info >= (3, 10) and USE_PYDANTIC:  # this doesn't work without pydantic type checking
+
+            def run_literal_pipe_test():
+                from typing import Literal
+
+                @csp.node
+                def node_with_literal_pipe(x: ts[int], choice: Literal["a", "b", "c"] | None | int) -> ts[str]:
+                    if csp.ticked(x):
+                        return str(choice) if choice is not None else "none"
+
+                @csp.graph
+                def graph_with_literal_pipe(choice: Literal["a", "b", "c"] | None | int) -> ts[str]:
+                    return csp.const(str(choice) if choice is not None else "none")
+
+                @csp.node
+                def dummy_node(x: ts["T"]):  # to avoid pruning
+                    if csp.ticked(x):
+                        ...
+
+                def graph():
+                    # These should work - valid literal values or None or int
+                    dummy_node(node_with_literal_pipe(csp.const(10), "a"))
+                    dummy_node(node_with_literal_pipe(csp.const(10), "b"))
+                    dummy_node(node_with_literal_pipe(csp.const(10), "c"))
+                    dummy_node(node_with_literal_pipe(csp.const(10), None))
+                    dummy_node(node_with_literal_pipe(csp.const(10), 12))
+
+                    graph_with_literal_pipe("a")
+                    graph_with_literal_pipe("b")
+                    graph_with_literal_pipe("c")
+                    graph_with_literal_pipe(None)
+                    graph_with_literal_pipe(12)
+
+                    msg = "(?s)2 validation errors for node_with_literal_pipe.*choice.*"
+                    with self.assertRaisesRegex(TypeError, msg):
+                        dummy_node(node_with_literal_pipe(csp.const(10), "d"))
+
+                csp.run(graph, starttime=datetime(2020, 2, 7, 9), endtime=datetime(2020, 2, 7, 9, 1))
+
+                # Test direct graph building
+                csp.build_graph(graph_with_literal_pipe, "a")
+                csp.build_graph(graph_with_literal_pipe, None)
+                csp.build_graph(graph_with_literal_pipe, 12)
+
+                # This should fail
+                msg = "(?s)2 validation errors for graph_with_literal_pipe.*choice.*"
+                with self.assertRaisesRegex(TypeError, msg):
+                    csp.build_graph(graph_with_literal_pipe, "d")
+
+            # Run the test
+            run_literal_pipe_test()
+
 
 if __name__ == "__main__":
     unittest.main()