From 0f8868639d21c29c338d38402a63c867236bdca1 Mon Sep 17 00:00:00 2001
From: Pascal Tomecek <pascal.tomecek@cubistsystematic.com>
Date: Mon, 14 Oct 2024 18:58:48 -0400
Subject: [PATCH] Add supported for Annotated to struct type.

Signed-off-by: Pascal Tomecek <pascal.tomecek@cubistsystematic.com>
---
 conda/dev-environment-unix.yml              |  1 +
 conda/dev-environment-win.yml               |  1 +
 csp/impl/types/container_type_normalizer.py |  3 +++
 csp/tests/impl/test_struct.py               | 17 +++++++++++++++++
 pyproject.toml                              |  2 ++
 5 files changed, 24 insertions(+)

diff --git a/conda/dev-environment-unix.yml b/conda/dev-environment-unix.yml
index 6ffafe73..e16b6f81 100644
--- a/conda/dev-environment-unix.yml
+++ b/conda/dev-environment-unix.yml
@@ -52,6 +52,7 @@ dependencies:
  - threadpoolctl
  - tornado
  - twine
+ - typing-extensions
  - unzip
  - wheel
  - zip
diff --git a/conda/dev-environment-win.yml b/conda/dev-environment-win.yml
index 227efc28..86e959b2 100644
--- a/conda/dev-environment-win.yml
+++ b/conda/dev-environment-win.yml
@@ -50,4 +50,5 @@ dependencies:
  - threadpoolctl
  - tornado
  - twine
+ - typing-extensions
  - wheel
diff --git a/csp/impl/types/container_type_normalizer.py b/csp/impl/types/container_type_normalizer.py
index ab7c94fd..9b4d5036 100644
--- a/csp/impl/types/container_type_normalizer.py
+++ b/csp/impl/types/container_type_normalizer.py
@@ -1,5 +1,6 @@
 import numpy
 import typing
+import typing_extensions
 
 import csp.typing
 from csp.impl.types.typing_utils import CspTypingUtils, FastList
@@ -70,6 +71,8 @@ def _convert_containers_to_typing_generic_meta(cls, typ, is_within_container):
 
     @classmethod
     def normalized_type_to_actual_python_type(cls, typ, level=0):
+        if isinstance(typ, typing_extensions._AnnotatedAlias):
+            typ = CspTypingUtils.get_origin(typ)
         if CspTypingUtils.is_generic_container(typ):
             if CspTypingUtils.get_origin(typ) is FastList and level == 0:
                 return [cls.normalized_type_to_actual_python_type(typ.__args__[0], level + 1), True]
diff --git a/csp/tests/impl/test_struct.py b/csp/tests/impl/test_struct.py
index ae9f2e45..1a5605b1 100644
--- a/csp/tests/impl/test_struct.py
+++ b/csp/tests/impl/test_struct.py
@@ -7,6 +7,7 @@
 import unittest
 from datetime import date, datetime, time, timedelta
 from typing import Dict, List, Set, Tuple
+from typing_extensions import Annotated
 
 import csp
 from csp.impl.struct import define_nested_struct, define_struct, defineNestedStruct, defineStruct
@@ -2943,6 +2944,22 @@ def test_dir(self):
         self.assertIn("__metadata__", dir_output)
         self.assertEqual(dir_output, sorted(dir_output))
 
+    def test_annotations(self):
+        class StructWithAnnotations(csp.Struct):
+            b: Annotated[float, "test"]
+            d: Annotated[Dict[str, Annotated[int, "test_int"]], "test_dict"]
+            s: str
+
+        self.assertEqual(
+            StructWithAnnotations.metadata(typed=True),
+            {
+                "b": Annotated[float, "test"],
+                "d": Annotated[Dict[str, Annotated[int, "test_int"]], "test_dict"],
+                "s": str,
+            },
+        )
+        self.assertEqual(StructWithAnnotations.metadata(typed=False), {"b": float, "d": dict, "s": str})
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/pyproject.toml b/pyproject.toml
index e9178160..f75cf325 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,6 +7,7 @@ requires = [
     "ruamel.yaml",
     "scikit-build",
     "setuptools>=69,<74",
+    "typing-extensions",
 ]
 build-backend="setuptools.build_meta"
 
@@ -29,6 +30,7 @@ dependencies = [
     "pytz",
     "ruamel.yaml",
     "sqlalchemy",
+    "typing-extensions",
 ]
 
 classifiers = [