Publicly expose is_nullable_numeric_dtype and get_data functions (#213)

Author: cbourjau

CHANGELOG.rst

@@ -7,6 +7,14 @@
 Changelog
 =========
 
+0.19.0 (2026-05-15)
+-------------------
+
+**New feature**
+
+- The :mod:`~ndonnx.extensions` module now publicly exposes the :func:`~ndonnx.extensions.is_nullable_numeric_dtype` and :func:`~ndonnx.extensions.get_data` functions.
+
+
 0.18.2 (2026-05-06)
 -------------------
 

ndonnx/extensions.py

@@ -10,9 +10,11 @@
 from typing_extensions import TypeIs, deprecated
 
 import ndonnx as ndx
-
-from . import _typed_array as tydx
-from ._typed_array.masked_onnx import TyMaArray
+import ndonnx._typed_array as tydx
+import ndonnx._typed_array.datetime
+import ndonnx._typed_array.funcs
+import ndonnx._typed_array.masked_onnx
+import ndonnx._typed_array.onnx
 
 SCALAR = TypeVar("SCALAR", int, float, str)
 
@@ -198,7 +200,7 @@ def make_nullable(
     null = None if null is None else null.copy()
 
     if null is None:
-        if isinstance(x._tyarray, TyMaArray):
+        if isinstance(x._tyarray, tydx.masked_onnx.TyMaArray):
             return x
         if isinstance(x._tyarray, tydx.onnx.TyArray):
             return ndx.Array._from_tyarray(
@@ -399,6 +401,12 @@ def is_nullable_dtype(dtype: ndx.DType, /) -> TypeIs[tydx.masked_onnx.DTypes]:
     return isinstance(dtype, tydx.masked_onnx.DTypes)
 
 
+def is_nullable_numeric_dtype(
+    dtype: ndx.DType, /
+) -> TypeIs[tydx.masked_onnx.NumericDTypes]:
+    return is_nullable_integer_dtype(dtype) or is_nullable_float_dtype(dtype)
+
+
 def is_nullable_integer_dtype(
     dtype: ndx.DType, /
 ) -> TypeIs[tydx.masked_onnx.IntegerDTypes]:
@@ -423,11 +431,13 @@ def is_time_unit(s: str, /) -> TypeIs[tydx.datetime.Unit]:
     "datetime_to_year_month_day",
     "fill_null",
     "get_mask",
+    "get_data",
     "is_float_dtype",
     "is_integer_dtype",
     "is_nullable_dtype",
     "is_nullable_float_dtype",
     "is_nullable_integer_dtype",
+    "is_nullable_numeric_dtype",
     "is_numeric_dtype",
     "is_onnx_dtype",
     "is_signed_integer_dtype",