Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Extend pandas with custom array types.
3"""
5from __future__ import annotations
7import inspect
8from typing import (
9 TYPE_CHECKING,
10 Any,
11 Self,
12 TypeVar,
13 cast,
14 overload,
15)
17import numpy as np
19from pandas._libs import missing as libmissing
20from pandas._libs.hashtable import object_hash
21from pandas._libs.properties import cache_readonly
22from pandas.errors import AbstractMethodError
23from pandas.util._decorators import set_module
25from pandas.core.dtypes.generic import (
26 ABCDataFrame,
27 ABCIndex,
28 ABCSeries,
29)
31if TYPE_CHECKING:
32 import matplotlib.units
34 from pandas._typing import (
35 DtypeObj,
36 Shape,
37 npt,
38 type_t,
39 )
41 from pandas import Index
42 from pandas.core.arrays import ExtensionArray
44 # To parameterize on same ExtensionDtype
45 ExtensionDtypeT = TypeVar("ExtensionDtypeT", bound="ExtensionDtype")
48@set_module("pandas.api.extensions")
49class ExtensionDtype:
50 """
51 A custom data type, to be paired with an ExtensionArray.
53 This enables support for third-party and custom dtypes within the
54 pandas ecosystem. By implementing this interface and pairing it with a custom
55 `ExtensionArray`, users can create rich data types that integrate cleanly
56 with pandas operations, such as grouping, joining, or aggregation.
58 Attributes
59 ----------
60 kind
61 na_value
62 name
63 names
64 type
66 Methods
67 -------
68 construct_array_type
69 construct_from_string
70 is_dtype
72 See Also
73 --------
74 extensions.register_extension_dtype: Register an ExtensionType
75 with pandas as class decorator.
76 extensions.ExtensionArray: Abstract base class for custom 1-D array types.
78 Notes
79 -----
80 The interface includes the following abstract methods that must
81 be implemented by subclasses:
83 * type
84 * name
85 * construct_array_type
87 The following attributes and methods influence the behavior of the dtype in
88 pandas operations
90 * _is_numeric
91 * _is_boolean
92 * _get_common_dtype
94 The `na_value` class attribute can be used to set the default NA value
95 for this type. :attr:`numpy.nan` is used by default.
97 ExtensionDtypes are required to be hashable. The base class provides
98 a default implementation, which relies on the ``_metadata`` class
99 attribute. ``_metadata`` should be a tuple containing the strings
100 that define your data type. For example, with ``PeriodDtype`` that's
101 the ``freq`` attribute.
103 **If you have a parametrized dtype you should set the ``_metadata``
104 class property**.
106 Ideally, the attributes in ``_metadata`` will match the
107 parameters to your ``ExtensionDtype.__init__`` (if any). If any of
108 the attributes in ``_metadata`` don't implement the standard
109 ``__eq__`` or ``__hash__``, the default implementations here will not
110 work.
112 Examples
113 --------
115 For interaction with Apache Arrow (pyarrow), a ``__from_arrow__`` method
116 can be implemented: this method receives a pyarrow Array or ChunkedArray
117 as only argument and is expected to return the appropriate pandas
118 ExtensionArray for this dtype and the passed values:
120 >>> import pyarrow
121 >>> from pandas.api.extensions import ExtensionArray
122 >>> class ExtensionDtype:
123 ... def __from_arrow__(
124 ... self, array: pyarrow.Array | pyarrow.ChunkedArray
125 ... ) -> ExtensionArray: ...
127 This class does not inherit from 'abc.ABCMeta' for performance reasons.
128 Methods and properties required by the interface raise
129 ``pandas.errors.AbstractMethodError`` and no ``register`` method is
130 provided for registering virtual subclasses.
131 """
133 _metadata: tuple[str, ...] = ()
135 def __str__(self) -> str:
136 return self.name
138 def __eq__(self, other: object) -> bool:
139 """
140 Check whether 'other' is equal to self.
142 By default, 'other' is considered equal if either
144 * it's a string matching 'self.name'.
145 * it's an instance of this type and all of the attributes
146 in ``self._metadata`` are equal between `self` and `other`.
148 Parameters
149 ----------
150 other : Any
152 Returns
153 -------
154 bool
155 """
156 if isinstance(other, str):
157 try:
158 other = self.construct_from_string(other)
159 except TypeError:
160 return False
161 if isinstance(other, type(self)):
162 return all(
163 getattr(self, attr) == getattr(other, attr) for attr in self._metadata
164 )
165 return False
167 def __hash__(self) -> int:
168 # different nan objects have different hashes
169 # we need to avoid that and thus use hash function with old behavior
170 return object_hash(tuple(getattr(self, attr) for attr in self._metadata))
172 def __ne__(self, other: object) -> bool:
173 return not self.__eq__(other)
175 @property
176 def na_value(self) -> object:
177 """
178 Default NA value to use for this type.
180 This is used in e.g. ExtensionArray.take. This should be the
181 user-facing "boxed" version of the NA value, not the physical NA value
182 for storage. e.g. for JSONArray, this is an empty dictionary.
183 """
184 return np.nan
186 @property
187 def type(self) -> type_t[Any]:
188 """
189 The scalar type for the array, e.g. ``int``
191 It's expected ``ExtensionArray[item]`` returns an instance
192 of ``ExtensionDtype.type`` for scalar ``item``, assuming
193 that value is valid (not NA). NA values do not need to be
194 instances of `type`.
195 """
196 raise AbstractMethodError(self)
198 @property
199 def kind(self) -> str:
200 """
201 A character code (one of 'biufcmMOSUV'), default 'O'
203 This should match the NumPy dtype used when the array is
204 converted to an ndarray, which is probably 'O' for object if
205 the extension type cannot be represented as a built-in NumPy
206 type.
208 See Also
209 --------
210 numpy.dtype.kind
211 """
212 return "O"
214 @property
215 def name(self) -> str:
216 """
217 A string identifying the data type.
219 Will be used for display in, e.g. ``Series.dtype``
220 """
221 raise AbstractMethodError(self)
223 @property
224 def names(self) -> list[str] | None:
225 """
226 Ordered list of field names, or None if there are no fields.
228 This is for compatibility with NumPy arrays, and may be removed in the
229 future.
230 """
231 return None
233 def construct_array_type(self) -> type_t[ExtensionArray]:
234 """
235 Return the array type associated with this dtype.
237 Returns
238 -------
239 type
240 """
241 raise AbstractMethodError(self)
243 def empty(self, shape: Shape) -> ExtensionArray:
244 """
245 Construct an ExtensionArray of this dtype with the given shape.
247 Analogous to numpy.empty.
249 Parameters
250 ----------
251 shape : int or tuple[int]
253 Returns
254 -------
255 ExtensionArray
256 """
257 cls = self.construct_array_type()
258 return cls._empty(shape, dtype=self)
260 @classmethod
261 def construct_from_string(cls, string: str) -> Self:
262 r"""
263 Construct this type from a string.
265 This is useful mainly for data types that accept parameters.
266 For example, a period dtype accepts a frequency parameter that
267 can be set as ``period[h]`` (where H means hourly frequency).
269 By default, in the abstract class, just the name of the type is
270 expected. But subclasses can overwrite this method to accept
271 parameters.
273 Parameters
274 ----------
275 string : str
276 The name of the type, for example ``category``.
278 Returns
279 -------
280 ExtensionDtype
281 Instance of the dtype.
283 Raises
284 ------
285 TypeError
286 If a class cannot be constructed from this 'string'.
288 Examples
289 --------
290 For extension dtypes with arguments the following may be an
291 adequate implementation.
293 >>> import re
294 >>> @classmethod
295 ... def construct_from_string(cls, string):
296 ... pattern = re.compile(r"^my_type\[(?P<arg_name>.+)\]$")
297 ... match = pattern.match(string)
298 ... if match:
299 ... return cls(**match.groupdict())
300 ... else:
301 ... raise TypeError(
302 ... f"Cannot construct a '{cls.__name__}' from '{string}'"
303 ... )
304 """
305 if not isinstance(string, str):
306 raise TypeError(
307 f"'construct_from_string' expects a string, got {type(string)}"
308 )
309 if not isinstance(cls.name, str):
310 # GH#46093 registered ExtensionDtype without a string `name`
311 raise TypeError(
312 f"Cannot construct a '{cls.__name__}' from a string because it "
313 "does not define a string 'name' attribute. ExtensionDtype "
314 "subclasses must set a class-level `name`."
315 )
316 # error: Non-overlapping equality check (left operand type: "str", right
317 # operand type: "Callable[[ExtensionDtype], str]") [comparison-overlap]
318 if string != cls.name:
319 raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")
320 return cls()
322 @classmethod
323 def is_dtype(cls, dtype: object) -> bool:
324 """
325 Check if we match 'dtype'.
327 Parameters
328 ----------
329 dtype : object
330 The object to check.
332 Returns
333 -------
334 bool
336 Notes
337 -----
338 The default implementation is True if
340 1. ``cls.construct_from_string(dtype)`` is an instance
341 of ``cls``.
342 2. ``dtype`` is an object and is an instance of ``cls``
343 3. ``dtype`` has a ``dtype`` attribute, and any of the above
344 conditions is true for ``dtype.dtype``.
345 """
346 dtype = getattr(dtype, "dtype", dtype)
348 if isinstance(dtype, (ABCSeries, ABCIndex, ABCDataFrame, np.dtype)):
349 # https://github.com/pandas-dev/pandas/issues/22960
350 # avoid passing data to `construct_from_string`. This could
351 # cause a FutureWarning from numpy about failing elementwise
352 # comparison from, e.g., comparing DataFrame == 'category'.
353 return False
354 elif dtype is None:
355 return False
356 elif isinstance(dtype, cls):
357 return True
358 if isinstance(dtype, str):
359 try:
360 return cls.construct_from_string(dtype) is not None
361 except TypeError:
362 return False
363 return False
365 @property
366 def _is_numeric(self) -> bool:
367 """
368 Whether columns with this dtype should be considered numeric.
370 By default ExtensionDtypes are assumed to be non-numeric.
371 They'll be excluded from operations that exclude non-numeric
372 columns, like (groupby) reductions, plotting, etc.
373 """
374 return False
376 @property
377 def _is_boolean(self) -> bool:
378 """
379 Whether this dtype should be considered boolean.
381 By default, ExtensionDtypes are assumed to be non-numeric.
382 Setting this to True will affect the behavior of several places,
383 e.g.
385 * is_bool
386 * boolean indexing
388 Returns
389 -------
390 bool
391 """
392 return False
394 def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
395 """
396 Return the common dtype, if one exists.
398 Used in `find_common_type` implementation. This is for example used
399 to determine the resulting dtype in a concat operation.
401 If no common dtype exists, return None (which gives the other dtypes
402 the chance to determine a common dtype). If all dtypes in the list
403 return None, then the common dtype will be "object" dtype (this means
404 it is never needed to return "object" dtype from this method itself).
406 Parameters
407 ----------
408 dtypes : list of dtypes
409 The dtypes for which to determine a common dtype. This is a list
410 of np.dtype or ExtensionDtype instances.
412 Returns
413 -------
414 Common dtype (np.dtype or ExtensionDtype) or None
415 """
416 if len(set(dtypes)) == 1:
417 # only itself
418 return self
419 else:
420 return None
422 @property
423 def _can_hold_na(self) -> bool:
424 """
425 Can arrays of this dtype hold NA values?
426 """
427 return True
429 @property
430 def _is_immutable(self) -> bool:
431 """
432 Can arrays with this dtype be modified with __setitem__? If not, return
433 True.
435 Immutable arrays are expected to raise TypeError on __setitem__ calls.
436 """
437 return False
439 @cache_readonly
440 def index_class(self) -> type_t[Index]:
441 """
442 The Index subclass to return from Index.__new__ when this dtype is
443 encountered.
444 """
445 from pandas import Index
447 return Index
449 @property
450 def _supports_2d(self) -> bool:
451 """
452 Do ExtensionArrays with this dtype support 2D arrays?
454 Historically ExtensionArrays were limited to 1D. By returning True here,
455 authors can indicate that their arrays support 2D instances. This can
456 improve performance in some cases, particularly operations with `axis=1`.
458 Arrays that support 2D values should:
460 - implement Array.reshape
461 - subclass the Dim2CompatTests in tests.extension.base
462 - _concat_same_type should support `axis` keyword
463 - _reduce and reductions should support `axis` keyword
464 """
465 return False
467 @property
468 def _can_fast_transpose(self) -> bool:
469 """
470 Is transposing an array with this dtype zero-copy?
472 Only relevant for cases where _supports_2d is True.
473 """
474 return False
476 @classmethod
477 def _get_plot_converter(
478 cls,
479 ) -> list[tuple[type_t, type_t[matplotlib.units.ConversionInterface]]]:
480 """
481 If dtype is plottable, return the type and converter to use for plotting.
483 By default only numeric ExtensionDtypes are plottable. In other cases, this
484 function just returns an empty list and ExtensionArrays of this dtype are
485 filtered out during plotting.
487 If the dtype is plottable, this function shall return a list of tuples of type
488 and converter pairs. These pairs are then registered to the matplotlib
489 converter registry, see
490 https://matplotlib.org/stable/api/units_api.html#module-matplotlib.units
491 The returned type is likely the same as ``cls.type``. The returned converter
492 should be a subclass of ``matplotlib.units.ConversionInterface``. If a dtype
493 supports multiple types, e.g. ``ArrowDtype``, then this function can return
494 multiple tuples inside the list, where each tuple is a pair of a type and
495 corresponding converter.
496 """
497 return []
500class StorageExtensionDtype(ExtensionDtype):
501 """ExtensionDtype that may be backed by more than one implementation."""
503 name: str
504 _metadata = ("storage",)
506 def __init__(self, storage: str) -> None:
507 self._storage = storage
509 def __repr__(self) -> str:
510 return f"{self.name}[{self.storage}]"
512 def __str__(self) -> str:
513 return self.name
515 def __eq__(self, other: object) -> bool:
516 if isinstance(other, str) and other == self.name:
517 return True
518 return super().__eq__(other)
520 def __hash__(self) -> int:
521 # custom __eq__ so have to override __hash__
522 return super().__hash__()
524 @property
525 def na_value(self) -> libmissing.NAType:
526 return libmissing.NA
528 @property
529 def storage(self) -> str:
530 return self._storage
533@set_module("pandas.api.extensions")
534def register_extension_dtype(cls: type_t[ExtensionDtypeT]) -> type_t[ExtensionDtypeT]:
535 """
536 Register an ExtensionType with pandas as class decorator.
538 This enables operations like ``.astype(name)`` for the name
539 of the ExtensionDtype.
541 Returns
542 -------
543 callable
544 A class decorator.
546 See Also
547 --------
548 api.extensions.ExtensionDtype : The base class for creating custom pandas
549 data types.
550 Series : One-dimensional array with axis labels.
551 DataFrame : Two-dimensional, size-mutable, potentially heterogeneous
552 tabular data.
554 Examples
555 --------
556 >>> from pandas.api.extensions import register_extension_dtype, ExtensionDtype
557 >>> @register_extension_dtype
558 ... class MyExtensionDtype(ExtensionDtype):
559 ... name = "myextension"
560 """
561 # GH#46093 identity check against the base property, so dtypes defining
562 # ``name`` as an instance-level property (e.g. DatetimeTZDtype) pass.
563 if inspect.getattr_static(cls, "name", None) is ExtensionDtype.__dict__["name"]:
564 raise TypeError(
565 f"Cannot register '{cls.__name__}' because it does not define a "
566 "string 'name' attribute. ExtensionDtype subclasses must set a "
567 "class-level `name`."
568 )
569 _registry.register(cls)
570 return cls
573class Registry:
574 """
575 Registry for dtype inference.
577 The registry allows one to map a string repr of an extension
578 dtype to an extension dtype. The string alias can be used in several
579 places, including
581 * Series and Index constructors
582 * :meth:`pandas.array`
583 * :meth:`pandas.Series.astype`
585 Multiple extension types can be registered.
586 These are tried in order.
587 """
589 def __init__(self) -> None:
590 self.dtypes: list[type_t[ExtensionDtype]] = []
592 def register(self, dtype: type_t[ExtensionDtype]) -> None:
593 """
594 Parameters
595 ----------
596 dtype : ExtensionDtype class
597 """
598 if not issubclass(dtype, ExtensionDtype):
599 raise ValueError("can only register pandas extension dtypes")
601 self.dtypes.append(dtype)
603 @overload
604 def find(self, dtype: type_t[ExtensionDtypeT]) -> type_t[ExtensionDtypeT]: ...
606 @overload
607 def find(self, dtype: ExtensionDtypeT) -> ExtensionDtypeT: ...
609 @overload
610 def find(self, dtype: str) -> ExtensionDtype | None: ...
612 @overload
613 def find(
614 self, dtype: npt.DTypeLike
615 ) -> type_t[ExtensionDtype] | ExtensionDtype | None: ...
617 def find(
618 self, dtype: type_t[ExtensionDtype] | ExtensionDtype | npt.DTypeLike
619 ) -> type_t[ExtensionDtype] | ExtensionDtype | None:
620 """
621 Parameters
622 ----------
623 dtype : ExtensionDtype class or instance or str or numpy dtype or python type
625 Returns
626 -------
627 return the first matching dtype, otherwise return None
628 """
629 if not isinstance(dtype, str):
630 dtype_type: type_t
631 if not isinstance(dtype, type):
632 dtype_type = type(dtype)
633 else:
634 dtype_type = dtype
635 if issubclass(dtype_type, ExtensionDtype):
636 # cast needed here as mypy doesn't know we have figured
637 # out it is an ExtensionDtype or type_t[ExtensionDtype]
638 return cast("ExtensionDtype | type_t[ExtensionDtype]", dtype)
640 return None
642 for dtype_type in self.dtypes:
643 try:
644 return dtype_type.construct_from_string(dtype)
645 except TypeError:
646 pass
648 return None
651_registry = Registry()