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
7from typing import (
8 TYPE_CHECKING,
9 Any,
10 Self,
11 TypeVar,
12 cast,
13 overload,
14)
16import numpy as np
18from pandas._libs import missing as libmissing
19from pandas._libs.hashtable import object_hash
20from pandas._libs.properties import cache_readonly
21from pandas.errors import AbstractMethodError
22from pandas.util._decorators import set_module
24from pandas.core.dtypes.generic import (
25 ABCDataFrame,
26 ABCIndex,
27 ABCSeries,
28)
30if TYPE_CHECKING:
31 from pandas._typing import (
32 DtypeObj,
33 Shape,
34 npt,
35 type_t,
36 )
38 from pandas import Index
39 from pandas.core.arrays import ExtensionArray
41 # To parameterize on same ExtensionDtype
42 ExtensionDtypeT = TypeVar("ExtensionDtypeT", bound="ExtensionDtype")
45@set_module("pandas.api.extensions")
46class ExtensionDtype:
47 """
48 A custom data type, to be paired with an ExtensionArray.
50 This enables support for third-party and custom dtypes within the
51 pandas ecosystem. By implementing this interface and pairing it with a custom
52 `ExtensionArray`, users can create rich data types that integrate cleanly
53 with pandas operations, such as grouping, joining, or aggregation.
55 See Also
56 --------
57 extensions.register_extension_dtype: Register an ExtensionType
58 with pandas as class decorator.
59 extensions.ExtensionArray: Abstract base class for custom 1-D array types.
61 Notes
62 -----
63 The interface includes the following abstract methods that must
64 be implemented by subclasses:
66 * type
67 * name
68 * construct_array_type
70 The following attributes and methods influence the behavior of the dtype in
71 pandas operations
73 * _is_numeric
74 * _is_boolean
75 * _get_common_dtype
77 The `na_value` class attribute can be used to set the default NA value
78 for this type. :attr:`numpy.nan` is used by default.
80 ExtensionDtypes are required to be hashable. The base class provides
81 a default implementation, which relies on the ``_metadata`` class
82 attribute. ``_metadata`` should be a tuple containing the strings
83 that define your data type. For example, with ``PeriodDtype`` that's
84 the ``freq`` attribute.
86 **If you have a parametrized dtype you should set the ``_metadata``
87 class property**.
89 Ideally, the attributes in ``_metadata`` will match the
90 parameters to your ``ExtensionDtype.__init__`` (if any). If any of
91 the attributes in ``_metadata`` don't implement the standard
92 ``__eq__`` or ``__hash__``, the default implementations here will not
93 work.
95 Examples
96 --------
98 For interaction with Apache Arrow (pyarrow), a ``__from_arrow__`` method
99 can be implemented: this method receives a pyarrow Array or ChunkedArray
100 as only argument and is expected to return the appropriate pandas
101 ExtensionArray for this dtype and the passed values:
103 >>> import pyarrow
104 >>> from pandas.api.extensions import ExtensionArray
105 >>> class ExtensionDtype:
106 ... def __from_arrow__(
107 ... self, array: pyarrow.Array | pyarrow.ChunkedArray
108 ... ) -> ExtensionArray: ...
110 This class does not inherit from 'abc.ABCMeta' for performance reasons.
111 Methods and properties required by the interface raise
112 ``pandas.errors.AbstractMethodError`` and no ``register`` method is
113 provided for registering virtual subclasses.
114 """
116 _metadata: tuple[str, ...] = ()
118 def __str__(self) -> str:
119 return self.name
121 def __eq__(self, other: object) -> bool:
122 """
123 Check whether 'other' is equal to self.
125 By default, 'other' is considered equal if either
127 * it's a string matching 'self.name'.
128 * it's an instance of this type and all of the attributes
129 in ``self._metadata`` are equal between `self` and `other`.
131 Parameters
132 ----------
133 other : Any
135 Returns
136 -------
137 bool
138 """
139 if isinstance(other, str):
140 try:
141 other = self.construct_from_string(other)
142 except TypeError:
143 return False
144 if isinstance(other, type(self)):
145 return all(
146 getattr(self, attr) == getattr(other, attr) for attr in self._metadata
147 )
148 return False
150 def __hash__(self) -> int:
151 # different nan objects have different hashes
152 # we need to avoid that and thus use hash function with old behavior
153 return object_hash(tuple(getattr(self, attr) for attr in self._metadata))
155 def __ne__(self, other: object) -> bool:
156 return not self.__eq__(other)
158 @property
159 def na_value(self) -> object:
160 """
161 Default NA value to use for this type.
163 This is used in e.g. ExtensionArray.take. This should be the
164 user-facing "boxed" version of the NA value, not the physical NA value
165 for storage. e.g. for JSONArray, this is an empty dictionary.
166 """
167 return np.nan
169 @property
170 def type(self) -> type_t[Any]:
171 """
172 The scalar type for the array, e.g. ``int``
174 It's expected ``ExtensionArray[item]`` returns an instance
175 of ``ExtensionDtype.type`` for scalar ``item``, assuming
176 that value is valid (not NA). NA values do not need to be
177 instances of `type`.
178 """
179 raise AbstractMethodError(self)
181 @property
182 def kind(self) -> str:
183 """
184 A character code (one of 'biufcmMOSUV'), default 'O'
186 This should match the NumPy dtype used when the array is
187 converted to an ndarray, which is probably 'O' for object if
188 the extension type cannot be represented as a built-in NumPy
189 type.
191 See Also
192 --------
193 numpy.dtype.kind
194 """
195 return "O"
197 @property
198 def name(self) -> str:
199 """
200 A string identifying the data type.
202 Will be used for display in, e.g. ``Series.dtype``
203 """
204 raise AbstractMethodError(self)
206 @property
207 def names(self) -> list[str] | None:
208 """
209 Ordered list of field names, or None if there are no fields.
211 This is for compatibility with NumPy arrays, and may be removed in the
212 future.
213 """
214 return None
216 def construct_array_type(self) -> type_t[ExtensionArray]:
217 """
218 Return the array type associated with this dtype.
220 Returns
221 -------
222 type
223 """
224 raise AbstractMethodError(self)
226 def empty(self, shape: Shape) -> ExtensionArray:
227 """
228 Construct an ExtensionArray of this dtype with the given shape.
230 Analogous to numpy.empty.
232 Parameters
233 ----------
234 shape : int or tuple[int]
236 Returns
237 -------
238 ExtensionArray
239 """
240 cls = self.construct_array_type()
241 return cls._empty(shape, dtype=self)
243 @classmethod
244 def construct_from_string(cls, string: str) -> Self:
245 r"""
246 Construct this type from a string.
248 This is useful mainly for data types that accept parameters.
249 For example, a period dtype accepts a frequency parameter that
250 can be set as ``period[h]`` (where H means hourly frequency).
252 By default, in the abstract class, just the name of the type is
253 expected. But subclasses can overwrite this method to accept
254 parameters.
256 Parameters
257 ----------
258 string : str
259 The name of the type, for example ``category``.
261 Returns
262 -------
263 ExtensionDtype
264 Instance of the dtype.
266 Raises
267 ------
268 TypeError
269 If a class cannot be constructed from this 'string'.
271 Examples
272 --------
273 For extension dtypes with arguments the following may be an
274 adequate implementation.
276 >>> import re
277 >>> @classmethod
278 ... def construct_from_string(cls, string):
279 ... pattern = re.compile(r"^my_type\[(?P<arg_name>.+)\]$")
280 ... match = pattern.match(string)
281 ... if match:
282 ... return cls(**match.groupdict())
283 ... else:
284 ... raise TypeError(
285 ... f"Cannot construct a '{cls.__name__}' from '{string}'"
286 ... )
287 """
288 if not isinstance(string, str):
289 raise TypeError(
290 f"'construct_from_string' expects a string, got {type(string)}"
291 )
292 # error: Non-overlapping equality check (left operand type: "str", right
293 # operand type: "Callable[[ExtensionDtype], str]") [comparison-overlap]
294 assert isinstance(cls.name, str), (cls, type(cls.name))
295 if string != cls.name:
296 raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")
297 return cls()
299 @classmethod
300 def is_dtype(cls, dtype: object) -> bool:
301 """
302 Check if we match 'dtype'.
304 Parameters
305 ----------
306 dtype : object
307 The object to check.
309 Returns
310 -------
311 bool
313 Notes
314 -----
315 The default implementation is True if
317 1. ``cls.construct_from_string(dtype)`` is an instance
318 of ``cls``.
319 2. ``dtype`` is an object and is an instance of ``cls``
320 3. ``dtype`` has a ``dtype`` attribute, and any of the above
321 conditions is true for ``dtype.dtype``.
322 """
323 dtype = getattr(dtype, "dtype", dtype)
325 if isinstance(dtype, (ABCSeries, ABCIndex, ABCDataFrame, np.dtype)):
326 # https://github.com/pandas-dev/pandas/issues/22960
327 # avoid passing data to `construct_from_string`. This could
328 # cause a FutureWarning from numpy about failing elementwise
329 # comparison from, e.g., comparing DataFrame == 'category'.
330 return False
331 elif dtype is None:
332 return False
333 elif isinstance(dtype, cls):
334 return True
335 if isinstance(dtype, str):
336 try:
337 return cls.construct_from_string(dtype) is not None
338 except TypeError:
339 return False
340 return False
342 @property
343 def _is_numeric(self) -> bool:
344 """
345 Whether columns with this dtype should be considered numeric.
347 By default ExtensionDtypes are assumed to be non-numeric.
348 They'll be excluded from operations that exclude non-numeric
349 columns, like (groupby) reductions, plotting, etc.
350 """
351 return False
353 @property
354 def _is_boolean(self) -> bool:
355 """
356 Whether this dtype should be considered boolean.
358 By default, ExtensionDtypes are assumed to be non-numeric.
359 Setting this to True will affect the behavior of several places,
360 e.g.
362 * is_bool
363 * boolean indexing
365 Returns
366 -------
367 bool
368 """
369 return False
371 def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
372 """
373 Return the common dtype, if one exists.
375 Used in `find_common_type` implementation. This is for example used
376 to determine the resulting dtype in a concat operation.
378 If no common dtype exists, return None (which gives the other dtypes
379 the chance to determine a common dtype). If all dtypes in the list
380 return None, then the common dtype will be "object" dtype (this means
381 it is never needed to return "object" dtype from this method itself).
383 Parameters
384 ----------
385 dtypes : list of dtypes
386 The dtypes for which to determine a common dtype. This is a list
387 of np.dtype or ExtensionDtype instances.
389 Returns
390 -------
391 Common dtype (np.dtype or ExtensionDtype) or None
392 """
393 if len(set(dtypes)) == 1:
394 # only itself
395 return self
396 else:
397 return None
399 @property
400 def _can_hold_na(self) -> bool:
401 """
402 Can arrays of this dtype hold NA values?
403 """
404 return True
406 @property
407 def _is_immutable(self) -> bool:
408 """
409 Can arrays with this dtype be modified with __setitem__? If not, return
410 True.
412 Immutable arrays are expected to raise TypeError on __setitem__ calls.
413 """
414 return False
416 @cache_readonly
417 def index_class(self) -> type_t[Index]:
418 """
419 The Index subclass to return from Index.__new__ when this dtype is
420 encountered.
421 """
422 from pandas import Index
424 return Index
426 @property
427 def _supports_2d(self) -> bool:
428 """
429 Do ExtensionArrays with this dtype support 2D arrays?
431 Historically ExtensionArrays were limited to 1D. By returning True here,
432 authors can indicate that their arrays support 2D instances. This can
433 improve performance in some cases, particularly operations with `axis=1`.
435 Arrays that support 2D values should:
437 - implement Array.reshape
438 - subclass the Dim2CompatTests in tests.extension.base
439 - _concat_same_type should support `axis` keyword
440 - _reduce and reductions should support `axis` keyword
441 """
442 return False
444 @property
445 def _can_fast_transpose(self) -> bool:
446 """
447 Is transposing an array with this dtype zero-copy?
449 Only relevant for cases where _supports_2d is True.
450 """
451 return False
454class StorageExtensionDtype(ExtensionDtype):
455 """ExtensionDtype that may be backed by more than one implementation."""
457 name: str
458 _metadata = ("storage",)
460 def __init__(self, storage: str) -> None:
461 self._storage = storage
463 def __repr__(self) -> str:
464 return f"{self.name}[{self.storage}]"
466 def __str__(self) -> str:
467 return self.name
469 def __eq__(self, other: object) -> bool:
470 if isinstance(other, str) and other == self.name:
471 return True
472 return super().__eq__(other)
474 def __hash__(self) -> int:
475 # custom __eq__ so have to override __hash__
476 return super().__hash__()
478 @property
479 def na_value(self) -> libmissing.NAType:
480 return libmissing.NA
482 @property
483 def storage(self) -> str:
484 return self._storage
487@set_module("pandas.api.extensions")
488def register_extension_dtype(cls: type_t[ExtensionDtypeT]) -> type_t[ExtensionDtypeT]:
489 """
490 Register an ExtensionType with pandas as class decorator.
492 This enables operations like ``.astype(name)`` for the name
493 of the ExtensionDtype.
495 Returns
496 -------
497 callable
498 A class decorator.
500 See Also
501 --------
502 api.extensions.ExtensionDtype : The base class for creating custom pandas
503 data types.
504 Series : One-dimensional array with axis labels.
505 DataFrame : Two-dimensional, size-mutable, potentially heterogeneous
506 tabular data.
508 Examples
509 --------
510 >>> from pandas.api.extensions import register_extension_dtype, ExtensionDtype
511 >>> @register_extension_dtype
512 ... class MyExtensionDtype(ExtensionDtype):
513 ... name = "myextension"
514 """
515 _registry.register(cls)
516 return cls
519class Registry:
520 """
521 Registry for dtype inference.
523 The registry allows one to map a string repr of an extension
524 dtype to an extension dtype. The string alias can be used in several
525 places, including
527 * Series and Index constructors
528 * :meth:`pandas.array`
529 * :meth:`pandas.Series.astype`
531 Multiple extension types can be registered.
532 These are tried in order.
533 """
535 def __init__(self) -> None:
536 self.dtypes: list[type_t[ExtensionDtype]] = []
538 def register(self, dtype: type_t[ExtensionDtype]) -> None:
539 """
540 Parameters
541 ----------
542 dtype : ExtensionDtype class
543 """
544 if not issubclass(dtype, ExtensionDtype):
545 raise ValueError("can only register pandas extension dtypes")
547 self.dtypes.append(dtype)
549 @overload
550 def find(self, dtype: type_t[ExtensionDtypeT]) -> type_t[ExtensionDtypeT]: ...
552 @overload
553 def find(self, dtype: ExtensionDtypeT) -> ExtensionDtypeT: ...
555 @overload
556 def find(self, dtype: str) -> ExtensionDtype | None: ...
558 @overload
559 def find(
560 self, dtype: npt.DTypeLike
561 ) -> type_t[ExtensionDtype] | ExtensionDtype | None: ...
563 def find(
564 self, dtype: type_t[ExtensionDtype] | ExtensionDtype | npt.DTypeLike
565 ) -> type_t[ExtensionDtype] | ExtensionDtype | None:
566 """
567 Parameters
568 ----------
569 dtype : ExtensionDtype class or instance or str or numpy dtype or python type
571 Returns
572 -------
573 return the first matching dtype, otherwise return None
574 """
575 if not isinstance(dtype, str):
576 dtype_type: type_t
577 if not isinstance(dtype, type):
578 dtype_type = type(dtype)
579 else:
580 dtype_type = dtype
581 if issubclass(dtype_type, ExtensionDtype):
582 # cast needed here as mypy doesn't know we have figured
583 # out it is an ExtensionDtype or type_t[ExtensionDtype]
584 return cast("ExtensionDtype | type_t[ExtensionDtype]", dtype)
586 return None
588 for dtype_type in self.dtypes:
589 try:
590 return dtype_type.construct_from_string(dtype)
591 except TypeError:
592 pass
594 return None
597_registry = Registry()