From a1a440e6ad32df790d7933d3e0fb8d3da3c39579 Mon Sep 17 00:00:00 2001
From: sneakers-the-rat <sneakers-the-rat@protonmail.com>
Date: Fri, 17 May 2024 18:05:36 -0700
Subject: [PATCH] handle and test complex

---
 src/numpydantic/maps.py   | 8 +++++++-
 src/numpydantic/schema.py | 6 ++++--
 tests/test_ndarray.py     | 6 +++++-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/numpydantic/maps.py b/src/numpydantic/maps.py
index 7de60c0..4db4b2d 100644
--- a/src/numpydantic/maps.py
+++ b/src/numpydantic/maps.py
@@ -58,5 +58,11 @@ flat_to_nptyping = {
 }
 """Map from NWB-style flat dtypes to nptyping types"""
 
-python_to_nptyping = {float: dt.Float, str: dt.String, int: dt.Int, bool: dt.Bool}
+python_to_nptyping = {
+    float: dt.Float,
+    str: dt.String,
+    int: dt.Int,
+    bool: dt.Bool,
+    complex: dt.Complex,
+}
 """Map from python types to nptyping types"""
diff --git a/src/numpydantic/schema.py b/src/numpydantic/schema.py
index b7733c6..5b22826 100644
--- a/src/numpydantic/schema.py
+++ b/src/numpydantic/schema.py
@@ -41,7 +41,6 @@ def _numeric_dtype(dtype: DtypeType, _handler: _handler_type) -> CoreSchema:
 
 def _lol_dtype(dtype: DtypeType, _handler: _handler_type) -> CoreSchema:
     """Get the innermost dtype schema to use in the generated pydantic schema"""
-
     if isinstance(dtype, nptyping.structure.StructureMeta):  # pragma: no cover
         raise NotImplementedError("Structured dtypes are currently unsupported")
 
@@ -63,7 +62,10 @@ def _lol_dtype(dtype: DtypeType, _handler: _handler_type) -> CoreSchema:
     else:
         try:
             python_type = np_to_python[dtype]
-        except KeyError as e:
+        except KeyError as e:  # pragma: no cover
+            # this should pretty much only happen in downstream/3rd-party interfaces
+            # that use interface-specific types. those need to provide mappings back
+            # to base python types (making this more streamlined is TODO)
             if dtype in np_to_python.values():
                 # it's already a python type
                 python_type = dtype
diff --git a/tests/test_ndarray.py b/tests/test_ndarray.py
index 88fe55b..f39dc8d 100644
--- a/tests/test_ndarray.py
+++ b/tests/test_ndarray.py
@@ -169,6 +169,7 @@ def test_json_schema_dtype_single(dtype, array_model):
         (int, "integer"),
         (float, "number"),
         (bool, "boolean"),
+        (complex, "any"),
     ],
 )
 def test_json_schema_dtype_builtin(dtype, expected, array_model):
@@ -179,7 +180,10 @@ def test_json_schema_dtype_builtin(dtype, expected, array_model):
     model = array_model(dtype=dtype)
     schema = model.model_json_schema()
     inner_type = schema["properties"]["array"]["items"]["items"]
-    assert inner_type["type"] == expected
+    if expected == "any":
+        assert inner_type == {}
+    else:
+        assert inner_type["type"] == expected
 
 
 @pytest.mark.skip("Not implemented yet")