From 43187bfa3ec6f722f4e5a0cbb304bae081b8d830 Mon Sep 17 00:00:00 2001
From: sneakers-the-rat <sneakers-the-rat@protonmail.com>
Date: Fri, 19 Jul 2024 20:50:48 -0700
Subject: [PATCH] add numpy type to language element types

---
 nwb_linkml/src/nwb_linkml/lang_elements.py | 35 ++++++++++++++--------
 nwb_linkml/src/nwb_linkml/maps/__init__.py |  3 +-
 nwb_linkml/src/nwb_linkml/maps/dtype.py    | 27 +++++++++++++++++
 3 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/nwb_linkml/src/nwb_linkml/lang_elements.py b/nwb_linkml/src/nwb_linkml/lang_elements.py
index d038c7e..b4f8fa6 100644
--- a/nwb_linkml/src/nwb_linkml/lang_elements.py
+++ b/nwb_linkml/src/nwb_linkml/lang_elements.py
@@ -2,7 +2,7 @@
 Language elements in nwb schema language that have a fixed, alternative representation
 in LinkML. These are exported as an nwb.language.yml file along with every generated namespace
 """
-
+from typing import List
 from linkml_runtime.linkml_model import (
     ClassDefinition,
     EnumDefinition,
@@ -12,7 +12,7 @@ from linkml_runtime.linkml_model import (
     TypeDefinition,
 )
 
-from nwb_linkml.maps import flat_to_linkml
+from nwb_linkml.maps import flat_to_linkml, flat_to_np
 from nwb_schema_language.datamodel.nwb_schema_pydantic import FlatDtype as FlatDtype_source
 
 FlatDType = EnumDefinition(
@@ -20,19 +20,28 @@ FlatDType = EnumDefinition(
     permissible_values=[PermissibleValue(p) for p in FlatDtype_source.__members__],
 )
 
-DTypeTypes = []
-for nwbtype, linkmltype in flat_to_linkml.items():
-    # skip the dtypes that are the same as the builtin linkml types (which should already exist)
-    # to avoid a recursion error
-    if linkmltype == nwbtype:
-        continue
 
-    amin = None
-    if nwbtype.startswith("uint"):
-        amin = 0
+def _make_dtypes() -> List[TypeDefinition]:
+    DTypeTypes = []
+    for nwbtype, linkmltype in flat_to_linkml.items():
+        # skip the dtypes that are the same as the builtin linkml types (which should already exist)
+        # to avoid a recursion error
+        if linkmltype == nwbtype:
+            continue
 
-    atype = TypeDefinition(name=nwbtype, minimum_value=amin, typeof=linkmltype)
-    DTypeTypes.append(atype)
+        amin = None
+        if nwbtype.startswith("uint"):
+            amin = 0
+
+        np_type = flat_to_np[nwbtype]
+
+        repr_string = f'np.{np_type.__name__}' if np_type.__module__ == 'numpy' else None
+
+        atype = TypeDefinition(name=nwbtype, minimum_value=amin, typeof=linkmltype, repr=repr_string)
+        DTypeTypes.append(atype)
+    return DTypeTypes
+
+DTypeTypes = _make_dtypes()
 
 AnyType = ClassDefinition(
     name="AnyType",
diff --git a/nwb_linkml/src/nwb_linkml/maps/__init__.py b/nwb_linkml/src/nwb_linkml/maps/__init__.py
index dac739d..b062d20 100644
--- a/nwb_linkml/src/nwb_linkml/maps/__init__.py
+++ b/nwb_linkml/src/nwb_linkml/maps/__init__.py
@@ -2,7 +2,7 @@
 Mapping from one domain to another
 """
 
-from nwb_linkml.maps.dtype import flat_to_linkml, flat_to_nptyping
+from nwb_linkml.maps.dtype import flat_to_linkml, flat_to_nptyping, flat_to_np
 from nwb_linkml.maps.map import Map
 from nwb_linkml.maps.postload import MAP_HDMF_DATATYPE_DEF, MAP_HDMF_DATATYPE_INC
 from nwb_linkml.maps.quantity import QUANTITY_MAP
@@ -13,5 +13,6 @@ __all__ = [
     "QUANTITY_MAP",
     "Map",
     "flat_to_linkml",
+    "flat_to_np",
     "flat_to_nptyping",
 ]
diff --git a/nwb_linkml/src/nwb_linkml/maps/dtype.py b/nwb_linkml/src/nwb_linkml/maps/dtype.py
index 2c4d2c1..9a7756f 100644
--- a/nwb_linkml/src/nwb_linkml/maps/dtype.py
+++ b/nwb_linkml/src/nwb_linkml/maps/dtype.py
@@ -69,6 +69,33 @@ flat_to_nptyping = {
     "object": "Object",
 }
 
+flat_to_np = {
+    "float": float,
+    "float32": np.float32,
+    "double": np.double,
+    "float64": np.float64,
+    "long": np.longlong,
+    "int64": np.int64,
+    "int": int,
+    "int32": np.int32,
+    "int16": np.int16,
+    "short": np.short,
+    "int8": np.int8,
+    "uint": np.uint,
+    "uint32": np.uint32,
+    "uint16": np.uint16,
+    "uint8": np.uint8,
+    "uint64": np.uint64,
+    "numeric": np.number,
+    "text": str,
+    "utf": str,
+    "utf8": str,
+    "utf_8": str,
+    "ascii": str,
+    "bool": bool,
+    "isodatetime": np.datetime64,
+}
+
 np_to_python = {
     Any: Any,
     np.number: float,