more OOP

cuvec/_common.py

@@ -1,11 +1,83 @@
 """Common helpers for cuvec.{cpython,pybind11,swig} modules."""
 import array
+import re
+from abc import ABC, abstractmethod
+from typing import Any, Dict
 from typing import Sequence as Seq
 from typing import Union
 
+import numpy as np
+
 Shape = Union[Seq[int], int]
 # u: non-standard np.dype('S2'); l/L: inconsistent between `array` and `numpy`
 typecodes = ''.join(i for i in array.typecodes if i not in "ulL")
+C_TYPES = {
+    "signed char": 'b',
+    "unsigned char": 'B',
+    "char": 'c',
+    "short": 'h',
+    "unsigned short": 'H',
+    "int": 'i',
+    "unsigned int": 'I',
+    "long long": 'q',
+    "unsigned long long": 'Q',
+    "float": 'f',
+    "double": 'd'} # yapf: disable
+
+
+class CVector(ABC):
+    """Thin wrapper around `CuVec<Type>`. Always takes ownership."""
+    vec_types: Dict[np.dtype, Any]
+    RE_CUVEC_TYPE: re.Pattern
+
+    def __init__(self, typechar: str):
+        self.typechar = typechar
+
+    @property
+    @abstractmethod
+    def shape(self) -> tuple:
+        pass
+
+    @property
+    @abstractmethod
+    def address(self) -> int:
+        pass
+
+    @property
+    def __array_interface__(self) -> Dict[str, Any]:
+        return {
+            'shape': self.shape, 'typestr': np.dtype(self.typechar).str,
+            'data': (self.address, False), 'version': 3}
+
+    @property
+    def __cuda_array_interface__(self) -> Dict[str, Any]:
+        return self.__array_interface__
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}('{self.typechar}', {self.shape})"
+
+    def __str__(self) -> str:
+        return f"{np.dtype(self.typechar)}{self.shape} at 0x{self.address:x}"
+
+    @classmethod
+    def zeros(cls, shape: Shape, dtype="float32"):
+        """Returns a new Vector of the specified shape and data type."""
+        return cls.vec_types[np.dtype(dtype)](shape)
+
+    @classmethod
+    def copy(cls, arr):
+        """Returns a new Vector with data copied from the specified `arr`."""
+        res = cls.zeros(arr.shape, arr.dtype)
+        np.asarray(res).flat = arr.flat
+        return res
+
+    @classmethod
+    def is_instance(cls, arr):
+        return isinstance(arr, cls) or type(arr).__name__ == cls.__name__
+
+    @classmethod
+    def is_raw_cuvec(cls, arr):
+        return cls.RE_CUVEC_TYPE.match(str(arr))
 
 
 def _generate_helpers(zeros, CuVec):

cuvec/pybind11.py

@@ -1,4 +1,8 @@
-"""Thin wrappers around `cuvec_pybind11` C++/CUDA module"""
+"""
+Thin wrappers around `cuvec_pybind11` C++/CUDA module
+
+A pybind11-driven equivalent of the CPython Extension API-driven `cpython.py`
+"""
 import logging
 import re
 from collections.abc import Sequence
@@ -9,45 +13,32 @@
 import numpy as np
 
 from . import cuvec_pybind11 as cu  # type: ignore [attr-defined] # yapf: disable
-from ._common import Shape, _generate_helpers, typecodes
+from ._common import C_TYPES, CVector, Shape, _generate_helpers, typecodes
 
 log = logging.getLogger(__name__)
-RE_NDCUVEC_TYPE = r"<.*NDCuVec_(.) object at 0x\w+>"
-NDCUVEC_TYPES = {
-    "signed char": 'b',
-    "unsigned char": 'B',
-    "char": 'c',
-    "short": 'h',
-    "unsigned short": 'H',
-    "int": 'i',
-    "unsigned int": 'I',
-    "long long": 'q',
-    "unsigned long long": 'Q',
-    "float": 'f',
-    "double": 'd'} # yapf: disable
+NDCUVEC_TYPES = dict(C_TYPES)
 if hasattr(cu, 'NDCuVec_e'):
     typecodes += 'e'
     NDCUVEC_TYPES["__half"] = 'e'
 
 
-class Pybind11Vector:
+class Pybind11Vector(CVector):
+    RE_CUVEC_TYPE = re.compile(r"<.*NDCuVec_(.) object at 0x\w+>")
+
     def __init__(self, typechar: str, shape: Shape, cuvec=None):
         """
-        Thin wrapper around `NDCuVec<Type>`. Always takes ownership.
         Args:
-          typechar(str)
+          typechar(char)
           shape(tuple(int))
           cuvec(NDCuVec<Type>): if given, `typechar` and `shape` are ignored
         """
-        if cuvec is not None:
-            assert is_raw_cuvec(cuvec)
-            self.typechar = re.match(RE_NDCUVEC_TYPE, str(cuvec)).group(1) # type: ignore
-            self.cuvec = cuvec
-            return
-
-        self.typechar = typechar
-        shape = cu.Shape(shape if isinstance(shape, Sequence) else (shape,))
-        self.cuvec = getattr(cu, f'NDCuVec_{typechar}')(shape)
+        if cuvec is None:
+            shape = cu.Shape(shape if isinstance(shape, Sequence) else (shape,))
+            cuvec = getattr(cu, f'NDCuVec_{typechar}')(shape)
+        else:
+            typechar = self.is_raw_cuvec(cuvec).group(1)
+        self.cuvec = cuvec
+        super().__init__(typechar)
 
     @property
     def shape(self) -> tuple:
@@ -57,48 +48,8 @@ def shape(self) -> tuple:
     def address(self) -> int:
         return self.cuvec.address()
 
-    @property
-    def __array_interface__(self) -> Dict[str, Any]:
-        return {
-            'shape': self.shape, 'typestr': np.dtype(self.typechar).str,
-            'data': (self.address, False), 'version': 3}
-
-    @property
-    def __cuda_array_interface__(self) -> Dict[str, Any]:
-        return self.__array_interface__
-
-    def __repr__(self) -> str:
-        return f"{type(self).__name__}('{self.typechar}', {self.shape})"
-
-    def __str__(self) -> str:
-        return f"{np.dtype(self.typechar)}{self.shape} at 0x{self.address:x}"
-
-
-vec_types = {np.dtype(c): partial(Pybind11Vector, c) for c in typecodes}
-
-
-def cu_zeros(shape: Shape, dtype="float32"):
-    """
-    Returns a new `Pybind11Vector` of the specified shape and data type.
-    """
-    return vec_types[np.dtype(dtype)](shape)
-
-
-def cu_copy(arr):
-    """
-    Returns a new `Pybind11Vector` with data copied from the specified `arr`.
-    """
-    res = cu_zeros(arr.shape, arr.dtype)
-    np.asarray(res).flat = arr.flat
-    return res
-
-
-def is_raw_cuvec(arr):
-    return re.match(RE_NDCUVEC_TYPE, str(arr))
-
 
-def is_raw_pyvec(arr):
-    return isinstance(arr, Pybind11Vector) or type(arr).__name__ == "Pybind11Vector"
+Pybind11Vector.vec_types = {np.dtype(c): partial(Pybind11Vector, c) for c in typecodes}
 
 
 class CuVec(np.ndarray):
@@ -108,7 +59,7 @@ class CuVec(np.ndarray):
     """
     def __new__(cls, arr):
         """arr: `cuvec.pybind11.CuVec`, raw `Pybind11Vector`, or `numpy.ndarray`"""
-        if is_raw_pyvec(arr):
+        if Pybind11Vector.is_instance(arr):
             log.debug("wrap pyraw %s", type(arr))
             obj = np.asarray(arr).view(cls)
             obj.pyvec = arr
@@ -146,7 +97,7 @@ def zeros(shape: Shape, dtype="float32") -> CuVec:
     Returns a `cuvec.pybind11.CuVec` view of a new `numpy.ndarray`
     of the specified shape and data type (`cuvec` equivalent of `numpy.zeros`).
     """
-    return CuVec(cu_zeros(shape, dtype))
+    return CuVec(Pybind11Vector.zeros(shape, dtype))
 
 
 ones, zeros_like, ones_like = _generate_helpers(zeros, CuVec)
@@ -158,7 +109,7 @@ def copy(arr) -> CuVec:
     with data copied from the specified `arr`
     (`cuvec` equivalent of `numpy.copy`).
     """
-    return CuVec(cu_copy(arr))
+    return CuVec(Pybind11Vector.copy(arr))
 
 
 def asarray(arr, dtype=None, order=None, ownership: str = 'warning') -> CuVec:
@@ -178,13 +129,13 @@ def asarray(arr, dtype=None, order=None, ownership: str = 'warning') -> CuVec:
         NB: `asarray()`/`retarray()` are safe if the raw cuvec was created in C++, e.g.:
         >>> res = retarray(some_pybind11_api_func(..., output=None))
     """
-    if is_raw_cuvec(arr):
+    if Pybind11Vector.is_raw_cuvec(arr):
         ownership = ownership.lower()
         if ownership in {'critical', 'fatal', 'error'}:
             raise IOError("Can't take ownership of existing cuvec (would create dangling ptr)")
         getattr(log, ownership)("taking ownership")
         arr = Pybind11Vector('', (), arr)
-    if not isinstance(arr, np.ndarray) and is_raw_pyvec(arr):
+    if not isinstance(arr, np.ndarray) and Pybind11Vector.is_instance(arr):
         res = CuVec(arr)
         if dtype is None or res.dtype == np.dtype(dtype):
             return CuVec(np.asanyarray(res, order=order))