complex: Add support for NumPy arrays of complex integers

docs/conf.py

@@ -56,7 +56,11 @@ def skip_aliases(app, what, name, obj, skip, options):
     """Skip documentation for classes that are exported from multiple modules."""
     # For names that are defined in a private sub-module and aliased into a
     # public package, hide the definition.
-    if name.startswith("nitypes.time._") or name.startswith("nitypes.waveform._"):
+    if (
+        name.startswith("nitypes.complex._")
+        or name.startswith("nitypes.time._")
+        or name.startswith("nitypes.waveform._")
+    ):
         skip = True
 
     return skip

src/nitypes/_arguments.py

@@ -6,7 +6,12 @@
 import numpy as np
 import numpy.typing as npt
 
-from nitypes._exceptions import invalid_arg_type, invalid_arg_value, unsupported_arg
+from nitypes._exceptions import (
+    invalid_arg_type,
+    invalid_arg_value,
+    unsupported_arg,
+    unsupported_dtype,
+)
 
 
 def arg_to_float(
@@ -144,14 +149,7 @@ def validate_dtype(dtype: npt.DTypeLike, supported_dtypes: tuple[npt.DTypeLike,
     if not isinstance(dtype, (type, np.dtype)):
         dtype = np.dtype(dtype)
     if not np.isdtype(dtype, supported_dtypes):
-        # Remove duplicate names because distinct types (e.g. int vs. long) may have the same name
-        # ("int32").
-        supported_dtype_names = {np.dtype(d).name: None for d in supported_dtypes}.keys()
-        raise TypeError(
-            "The requested data type is not supported.\n\n"
-            f"Data type: {np.dtype(dtype)}\n"
-            f"Supported data types: {', '.join(supported_dtype_names)}"
-        )
+        raise unsupported_dtype("requested data type", dtype, supported_dtypes)
 
 
 def validate_unsupported_arg(arg_description: str, value: object) -> None:

src/nitypes/_exceptions.py

@@ -3,6 +3,9 @@
 import reprlib
 import sys
 
+import numpy as np
+import numpy.typing as npt
+
 
 def add_note(exception: Exception, note: str) -> None:
     """Add a note to an exception.
@@ -66,6 +69,20 @@ def unsupported_arg(arg_description: str, value: object) -> ValueError:
     )
 
 
+def unsupported_dtype(
+    arg_description: str, dtype: npt.DTypeLike, supported_dtypes: tuple[npt.DTypeLike, ...]
+) -> TypeError:
+    """Create a TypeError for an unsupported dtype."""
+    # Remove duplicate names because distinct types (e.g. int vs. long) may have the same name
+    # ("int32").
+    supported_dtype_names = {np.dtype(d).name: None for d in supported_dtypes}.keys()
+    return TypeError(
+        f"The {arg_description} is not supported.\n\n"
+        f"Data type: {np.dtype(dtype)}\n"
+        f"Supported data types: {', '.join(supported_dtype_names)}"
+    )
+
+
 # English-specific hack. This is why we prefer "Key: value" for localizable errors. TODO: consider
 # moving the full strings into a string table instead of building them out of English noun phrases.
 def _a(noun: str) -> str:

src/nitypes/complex/__init__.py

@@ -0,0 +1,129 @@
+"""Complex number data types for NI Python APIs.
+
+================
+Complex Integers
+================
+
+Some NI driver APIs (such as NI-FGEN, NI-SCOPE, NI-RFSA, and NI-RFSG) use complex numbers to
+represent I/Q data. Python and NumPy have native support for complex floating-point numbers, but
+not complex integers, so the :mod:`nitypes.complex` submodule provides a NumPy representation of
+complex integers.
+
+:any:`ComplexInt32DType` is a NumPy structured data type object representing a complex integer with
+16-bit ``real`` and ``imag`` fields. This structured data type has the same memory layout as the
+``NIComplexI16`` C struct used by NI driver APIs.
+
+For more information about NumPy structured data types, see the
+:ref:`NumPy documentation on structured arrays <numpy:structured_arrays>`.
+
+.. note::
+   In ``NIComplexI16``, the number 16 refers to the number of bits in each field. In
+   :any:`ComplexInt32DType`, the number 32 refers to the total number of bits, following the
+   precedent set by NumPy's other complex types. For example, :any:`numpy.complex128` contains
+   64-bit ``real`` and ``imag`` fields.
+
+Constructing arrays of complex integers
+---------------------------------------
+
+You can construct an array of complex integers from a sequence of tuples using :func:`numpy.array`:
+
+>>> import numpy as np
+>>> np.array([(1, 2), (3, 4)], dtype=ComplexInt32DType)
+array([(1, 2), (3, 4)], dtype=[('real', '<i2'), ('imag', '<i2')])
+
+Likewise, you can construct an array of complex integer zeros using :func:`numpy.zeros`:
+
+>>> np.zeros(3, dtype=ComplexInt32DType)
+array([(0, 0), (0, 0), (0, 0)], dtype=[('real', '<i2'), ('imag', '<i2')])
+
+Indexing and slicing
+--------------------
+
+Indexing the array gives you a complex integer structured scalar:
+
+>>> x = np.array([(1, 2), (3, 4), (5, 6)], dtype=ComplexInt32DType)
+>>> x[0]
+np.void((1, 2), dtype=[('real', '<i2'), ('imag', '<i2')])
+>>> x[1]
+np.void((3, 4), dtype=[('real', '<i2'), ('imag', '<i2')])
+
+.. note:
+    NumPy displays :any:`numpy.void` because the :any:`ComplexInt32DType` structured data type has
+    a base type of :any:`numpy.void`. Using a different base type such as :any:`numpy.int32`
+    would have benefits, such as making it easier to convert array elements to/from
+    :any:`numpy.int32`, but it would also have drawbacks, such as making it harder to initialize
+    the array using a sequence of tuples.
+
+You can index a complex integer structured scalar to get the real and imaginary parts:
+
+>>> x[0][0]
+np.int16(1)
+>>> x[0][1]
+np.int16(2)
+
+You can also index by the field names ``real`` and ``imag``:
+
+>>> x[0]['real']
+np.int16(1)
+>>> x[0]['imag']
+np.int16(2)
+
+Or you can index the entire array by the field names ``real`` and ``imag``:
+
+>>> x['real']
+array([1, 3, 5], dtype=int16)
+>>> x['imag']
+array([2, 4, 6], dtype=int16)
+
+Arrays of complex integers support slicing and negative indices like any other array:
+
+>>> x[0:2]
+array([(1, 2), (3, 4)], dtype=[('real', '<i2'), ('imag', '<i2')])
+>>> x[1:]
+array([(3, 4), (5, 6)], dtype=[('real', '<i2'), ('imag', '<i2')])
+>>> x[-1]
+np.void((5, 6), dtype=[('real', '<i2'), ('imag', '<i2')])
+
+Conversion
+----------
+
+To convert a complex integer structured scalar to a tuple, use the :any:`numpy.ndarray.item`
+method:
+
+>>> x[0].item()
+(1, 2)
+>>> [y.item() for y in x]
+[(1, 2), (3, 4), (5, 6)]
+
+To convert NumPy arrays between between different complex number data types, use the
+:func:`convert_complex` function:
+
+>>> convert_complex(np.complex128, x)
+array([1.+2.j, 3.+4.j, 5.+6.j])
+>>> convert_complex(ComplexInt32DType, np.array([1.23+4.56j]))
+array([(1, 4)], dtype=[('real', '<i2'), ('imag', '<i2')])
+
+You can also use :func:`convert_complex` with NumPy scalars:
+
+>>> convert_complex(np.complex128, x[0])
+np.complex128(1+2j)
+>>> convert_complex(ComplexInt32DType, np.complex128(3+4j))
+np.void((3, 4), dtype=[('real', '<i2'), ('imag', '<i2')])
+
+.. note::
+    As of NumPy 2.2, shape typing is still under development, so its type hints do not reflect that
+    many operations coerce zero-dimensional arrays to :any:`numpy.generic`. The type hints for the
+    scalar overloads of :func:`convert_complex` follow this precedent and return an
+    :any:`numpy.ndarray`. This behavior may change in a future release.
+
+Mathematical operations
+-----------------------
+
+Structured arrays of complex integers do not support mathematical operations. Convert
+them to arrays of complex floating-point numbers before doing any sort of math or analysis.
+"""
+
+from nitypes.complex._conversion import convert_complex
+from nitypes.complex._dtypes import ComplexInt32Base, ComplexInt32DType
+
+__all__ = ["convert_complex", "ComplexInt32DType", "ComplexInt32Base"]

src/nitypes/complex/_conversion.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from typing import Any, TypeVar, cast, overload
+
+import numpy as np
+import numpy.typing as npt
+
+from nitypes._arguments import validate_dtype
+from nitypes._exceptions import unsupported_dtype
+from nitypes.complex._dtypes import ComplexInt32DType
+
+_Item_co = TypeVar("_Item_co", bound=Any)
+_ScalarType = TypeVar("_ScalarType", bound=np.generic)
+_Shape = TypeVar("_Shape", bound=tuple[int, ...])
+
+_COMPLEX_DTYPES = (
+    np.complex64,
+    np.complex128,
+    ComplexInt32DType,
+)
+
+_FIELD_DTYPE = {
+    np.dtype(np.complex64): np.float32,
+    np.dtype(np.complex128): np.float64,
+    ComplexInt32DType: np.int16,
+}
+
+
+@overload
+def convert_complex(
+    requested_dtype: type[_ScalarType] | np.dtype[_ScalarType],
+    value: np.ndarray[_Shape, Any],
+) -> np.ndarray[_Shape, np.dtype[_ScalarType]]: ...
+
+
+@overload
+def convert_complex(
+    requested_dtype: npt.DTypeLike, value: np.ndarray[_Shape, Any]
+) -> np.ndarray[_Shape, Any]: ...
+
+
+# https://numpy.org/doc/2.2/reference/typing.html#d-arrays
+# "While thus not strictly correct, all operations are that can potentially perform a 0D-array ->
+# scalar cast are currently annotated as exclusively returning an ndarray."
+@overload
+def convert_complex(
+    requested_dtype: type[_ScalarType] | np.dtype[_ScalarType],
+    value: np.generic[Any],
+) -> np.ndarray[tuple[()], np.dtype[_ScalarType]]: ...
+
+
+@overload
+def convert_complex(
+    requested_dtype: npt.DTypeLike,
+    value: np.generic[Any],
+) -> np.ndarray[tuple[()], Any]: ...
+
+
+def convert_complex(
+    requested_dtype: npt.DTypeLike, value: np.ndarray[_Shape, Any] | np.generic[Any]
+) -> np.ndarray[_Shape, Any]:
+    """Convert a NumPy array or scalar of complex numbers to the specified dtype.
+
+    Args:
+        requested_dtype: The NumPy data type to convert to. Supported data types:
+            :any:`numpy.complex64`, :any:`numpy.complex128`, :any:`ComplexInt32DType`.
+        value: The NumPy array or scalar to convert.
+
+    Returns:
+        The value converted to the specified dtype.
+    """
+    validate_dtype(requested_dtype, _COMPLEX_DTYPES)
+    if requested_dtype == value.dtype:
+        return cast(np.ndarray[_Shape, Any], value)
+    elif requested_dtype == ComplexInt32DType or value.dtype == ComplexInt32DType:
+        # ndarray.view on scalars requires the source and destination types to have the same size,
+        # so reshape the scalar into an 1-element array before converting and index it afterwards.
+        # shape == () means this is either a scalar (np.generic) or a 0-dimension array, but mypy
+        # doesn't know that.
+        if value.shape == ():
+            return cast(
+                np.ndarray[_Shape, Any],
+                _convert_complexint32_array(requested_dtype, value.reshape(1))[0],
+            )
+        else:
+            return _convert_complexint32_array(
+                requested_dtype, cast(np.ndarray[_Shape, Any], value)
+            )
+    else:
+        return value.astype(requested_dtype)
+
+
+def _convert_complexint32_array(
+    requested_dtype: npt.DTypeLike | type[_ScalarType] | np.dtype[_ScalarType],
+    value: np.ndarray[_Shape, Any],
+) -> np.ndarray[_Shape, np.dtype[_ScalarType]]:
+    if not isinstance(requested_dtype, np.dtype):
+        requested_dtype = np.dtype(requested_dtype)
+
+    requested_field_dtype = _FIELD_DTYPE.get(requested_dtype)
+    if requested_field_dtype is None:
+        raise unsupported_dtype("requested data type", requested_dtype, _COMPLEX_DTYPES)
+
+    value_field_dtype = _FIELD_DTYPE.get(value.dtype)
+    if value_field_dtype is None:
+        raise unsupported_dtype("array data type", value.dtype, _COMPLEX_DTYPES)
+
+    return value.view(value_field_dtype).astype(requested_field_dtype).view(requested_dtype)

src/nitypes/complex/_dtypes.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+import numpy as np
+from typing_extensions import TypeAlias
+
+ComplexInt32Base: TypeAlias = np.void
+"""Type alias for the base type of :any:`ComplexInt32DType`, which is :any:`numpy.void`."""
+
+ComplexInt32DType = np.dtype((ComplexInt32Base, [("real", np.int16), ("imag", np.int16)]))
+"""NumPy structured data type for a complex integer with 16-bit ``real`` and ``imag`` fields."""

tests/unit/complex/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for the nitypes.complex package."""