Videre

linedance-app/venv/lib/python3.12/site-packages/numba/tests/pdlike_usecase.py

@@ -0,0 +1,306 @@
+"""
+Implementation of a minimal Pandas-like API.
+"""
+
+import numpy as np
+
+from numba.core import types, cgutils
+from numba.core.datamodel import models
+from numba.core.extending import (
+    typeof_impl, type_callable, register_model,
+    lower_builtin, box, unbox, NativeValue,
+    overload, overload_attribute, overload_method, make_attribute_wrapper)
+from numba.core.imputils import impl_ret_borrowed
+
+
+class Index(object):
+    """
+    A minimal pandas.Index-like object.
+    """
+
+    def __init__(self, data):
+        assert isinstance(data, np.ndarray)
+        assert data.ndim == 1
+        self._data = data
+
+    def __iter__(self):
+        return iter(self._data)
+
+    @property
+    def dtype(self):
+        return self._data.dtype
+
+    @property
+    def flags(self):
+        return self._data.flags
+
+
+class IndexType(types.Buffer):
+    """
+    The type class for Index objects.
+    """
+    array_priority = 1000
+
+    def __init__(self, dtype, layout, pyclass):
+        self.pyclass = pyclass
+        super(IndexType, self).__init__(dtype, 1, layout)
+
+    @property
+    def key(self):
+        return self.pyclass, self.dtype, self.layout
+
+    @property
+    def as_array(self):
+        return types.Array(self.dtype, 1, self.layout)
+
+    def copy(self, dtype=None, ndim=1, layout=None):
+        assert ndim == 1
+        if dtype is None:
+            dtype = self.dtype
+        layout = layout or self.layout
+        return type(self)(dtype, layout, self.pyclass)
+
+
+class Series(object):
+    """
+    A minimal pandas.Series-like object.
+    """
+
+    def __init__(self, data, index):
+        assert isinstance(data, np.ndarray)
+        assert isinstance(index, Index)
+        assert data.ndim == 1
+        self._values = data
+        self._index = index
+
+    def __iter__(self):
+        return iter(self._values)
+
+    @property
+    def dtype(self):
+        return self._values.dtype
+
+    @property
+    def flags(self):
+        return self._values.flags
+
+
+class SeriesType(types.ArrayCompatible):
+    """
+    The type class for Series objects.
+    """
+    array_priority = 1000
+
+    def __init__(self, dtype, index):
+        assert isinstance(index, IndexType)
+        self.dtype = dtype
+        self.index = index
+        self.values = types.Array(self.dtype, 1, 'C')
+        name = "series(%s, %s)" % (dtype, index)
+        super(SeriesType, self).__init__(name)
+
+    @property
+    def key(self):
+        return self.dtype, self.index
+
+    @property
+    def as_array(self):
+        return self.values
+
+    def copy(self, dtype=None, ndim=1, layout='C'):
+        assert ndim == 1
+        assert layout == 'C'
+        if dtype is None:
+            dtype = self.dtype
+        return type(self)(dtype, self.index)
+
+
+@typeof_impl.register(Index)
+def typeof_index(val, c):
+    arrty = typeof_impl(val._data, c)
+    assert arrty.ndim == 1
+    return IndexType(arrty.dtype, arrty.layout, type(val))
+
+@typeof_impl.register(Series)
+def typeof_series(val, c):
+    index = typeof_impl(val._index, c)
+    arrty = typeof_impl(val._values, c)
+    assert arrty.ndim == 1
+    assert arrty.layout == 'C'
+    return SeriesType(arrty.dtype, index)
+
+@type_callable('__array_wrap__')
+def type_array_wrap(context):
+    def typer(input_type, result):
+        if isinstance(input_type, (IndexType, SeriesType)):
+            return input_type.copy(dtype=result.dtype,
+                                   ndim=result.ndim,
+                                   layout=result.layout)
+
+    return typer
+
+@type_callable(Series)
+def type_series_constructor(context):
+    def typer(data, index):
+        if isinstance(index, IndexType) and isinstance(data, types.Array):
+            assert data.layout == 'C'
+            assert data.ndim == 1
+            return SeriesType(data.dtype, index)
+
+    return typer
+
+
+# Backend extensions for Index and Series
+
+@register_model(IndexType)
+class IndexModel(models.StructModel):
+    def __init__(self, dmm, fe_type):
+        members = [('data', fe_type.as_array)]
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+@register_model(SeriesType)
+class SeriesModel(models.StructModel):
+    def __init__(self, dmm, fe_type):
+        members = [
+            ('index', fe_type.index),
+            ('values', fe_type.as_array),
+            ]
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+make_attribute_wrapper(IndexType, 'data', '_data')
+make_attribute_wrapper(SeriesType, 'index', '_index')
+make_attribute_wrapper(SeriesType, 'values', '_values')
+
+def make_index(context, builder, typ, **kwargs):
+    return cgutils.create_struct_proxy(typ)(context, builder, **kwargs)
+
+def make_series(context, builder, typ, **kwargs):
+    return cgutils.create_struct_proxy(typ)(context, builder, **kwargs)
+
+@lower_builtin('__array__', IndexType)
+def index_as_array(context, builder, sig, args):
+    val = make_index(context, builder, sig.args[0], ref=args[0])
+    return val._get_ptr_by_name('data')
+
+@lower_builtin('__array__', SeriesType)
+def series_as_array(context, builder, sig, args):
+    val = make_series(context, builder, sig.args[0], ref=args[0])
+    return val._get_ptr_by_name('values')
+
+@lower_builtin('__array_wrap__', IndexType, types.Array)
+def index_wrap_array(context, builder, sig, args):
+    dest = make_index(context, builder, sig.return_type)
+    dest.data = args[1]
+    return impl_ret_borrowed(context, builder, sig.return_type, dest._getvalue())
+
+@lower_builtin('__array_wrap__', SeriesType, types.Array)
+def series_wrap_array(context, builder, sig, args):
+    src = make_series(context, builder, sig.args[0], value=args[0])
+    dest = make_series(context, builder, sig.return_type)
+    dest.values = args[1]
+    dest.index = src.index
+    return impl_ret_borrowed(context, builder, sig.return_type, dest._getvalue())
+
+@lower_builtin(Series, types.Array, IndexType)
+def pdseries_constructor(context, builder, sig, args):
+    data, index = args
+    series = make_series(context, builder, sig.return_type)
+    series.index = index
+    series.values = data
+    return impl_ret_borrowed(context, builder, sig.return_type, series._getvalue())
+
+
+@unbox(IndexType)
+def unbox_index(typ, obj, c):
+    """
+    Convert a Index object to a native structure.
+    """
+    data = c.pyapi.object_getattr_string(obj, "_data")
+    index = make_index(c.context, c.builder, typ)
+    index.data = c.unbox(typ.as_array, data).value
+
+    return NativeValue(index._getvalue())
+
+@unbox(SeriesType)
+def unbox_series(typ, obj, c):
+    """
+    Convert a Series object to a native structure.
+    """
+    index = c.pyapi.object_getattr_string(obj, "_index")
+    values = c.pyapi.object_getattr_string(obj, "_values")
+    series = make_series(c.context, c.builder, typ)
+    series.index = c.unbox(typ.index, index).value
+    series.values = c.unbox(typ.values, values).value
+
+    return NativeValue(series._getvalue())
+
+
+@box(IndexType)
+def box_index(typ, val, c):
+    """
+    Convert a native index structure to a Index object.
+    """
+    # First build a Numpy array object, then wrap it in a Index
+    index = make_index(c.context, c.builder, typ, value=val)
+    classobj = c.pyapi.unserialize(c.pyapi.serialize_object(typ.pyclass))
+    arrayobj = c.box(typ.as_array, index.data)
+    indexobj = c.pyapi.call_function_objargs(classobj, (arrayobj,))
+    return indexobj
+
+@box(SeriesType)
+def box_series(typ, val, c):
+    """
+    Convert a native series structure to a Series object.
+    """
+    series = make_series(c.context, c.builder, typ, value=val)
+    classobj = c.pyapi.unserialize(c.pyapi.serialize_object(Series))
+    indexobj = c.box(typ.index, series.index)
+    arrayobj = c.box(typ.as_array, series.values)
+    seriesobj = c.pyapi.call_function_objargs(classobj, (arrayobj, indexobj))
+    return seriesobj
+
+
+@overload_attribute(IndexType, 'is_monotonic_increasing')
+def index_is_monotonic_increasing(index):
+    """
+    Index.is_monotonic_increasing
+    """
+    def getter(index):
+        data = index._data
+        if len(data) == 0:
+            return True
+        u = data[0]
+        for v in data:
+            if v < u:
+                return False
+            v = u
+        return True
+
+    return getter
+
+@overload(len)
+def series_len(series):
+    """
+    len(Series)
+    """
+    if isinstance(series, SeriesType):
+        def len_impl(series):
+            return len(series._values)
+        return len_impl
+
+@overload_method(SeriesType, 'clip')
+def series_clip(series, lower, upper):
+    """
+    Series.clip(...)
+    """
+    def clip_impl(series, lower, upper):
+        data = series._values.copy()
+        for i in range(len(data)):
+            v = data[i]
+            if v < lower:
+                data[i] = lower
+            elif v > upper:
+                data[i] = upper
+        return Series(data, series._index)
+
+    return clip_impl