Fix typing for extension arrays and extension dtypes without isin and astype #40421

Dr-Irv · 2021-03-13T17:27:20Z

Same as #39501 by removing changes related to isin and the overloads for astype and return type for astype

Per discussion with @simonjayhawkins #39501 (comment)

… and return type

pandas/core/arrays/base.py

Dr-Irv · 2021-03-13T22:21:57Z

@Dr-Irv I think we still need to reduce the size of this PR, so need to start breaking non controversial bits off so we can start merging some of this great work (i.e. changing typvars, the getitem overloads and maybe privatising registry need further input from others)

OK, next commit will remove those three things, which were changing the TypeVar on DatetimeLikeScalar, removing the __getitem__ overloads and the changing of registry to _registry

simonjayhawkins · 2021-03-13T22:23:01Z

pandas/core/internals/blocks.py

-                # error: Invalid index type "List[Any]" for "ExtensionArray"; expected
-                # type "Union[int, slice, ndarray]"
-                return self.values[[loc]]  # type: ignore[index]
+                return self.values[[loc]]


This may not be a false positive. The EA interface does not specify (in the docstring) that a list is accepted for __getitem__, although maybe we have tests in the base extension tests that enforce this. will need to check.

Maybe the docs need to get updated, because this works:

>>> s=pd.array([1,2,3,4,5]) >>> s[[2,4]] <IntegerArray> [3, 5] Length: 2, dtype: Int64

sure for our internal EAs. it's 3rd party EAs that may not support this.

So I guess there are two choices:

Require 3rd party EAs to support it and update the docs

Remove the List[Any] that I put in, and just mark the mypy error on that one place

I will assume you prefer (2), and will remove that

will look in detail tomorrow.

Last commits do it using (2).

Require 3rd party EAs to support it and update the docs

This one (not necessarily for this PR)

One other thought on this. There was some comment about having the EA's include a 2D signature, but I think that can be done in a subclass that implements it. So right now, the signature for __getitem__() is:

def __getitem__( self, item: Union[int, slice, np.ndarray] ) -> Union[ExtensionArray, Any]:

The discussion here is (to do in another PR) to make it (and change the docs):

def __getitem__( self, item: Union[int, slice, np.ndarray, List[int]] ) -> Union[ExtensionArray, Any]:

Note the use of List[int] here, not List[Any] as I had it before.

If there is a 2D EA that wants to allow Tuple[int, ...], it can have it's own signature that widens the EA signature:

def __getitem__( self, item: Union[int, slice, np.ndarray, List[int], Tuple[int, ...]] ) -> Union[ExtensionArray, Any]:

But we don't have to include that signature with Tuple[int,...] in the base EA class.

pandas/io/formats/format.py

jbrockmendel · 2021-03-13T23:22:04Z

pandas/core/base.py

@@ -1073,7 +1069,7 @@ def unique(self):
        values = self._values

        if not isinstance(values, np.ndarray):
-            result = values.unique()
+            result: ArrayLike = values.unique()


we've got ExtensionArray in the namespace, can make the isinstance check above for EA

I've changed the isinstance check, but you still need to declare result as ArrayLike, because the code below that line can return the result as an np.ndarray . Here's the excerpt (as now written):

if isinstance(values, ExtensionArray): result: ArrayLike = values.unique() if self.dtype.kind in ["m", "M"] and isinstance(self, ABCSeries): # GH#31182 Series._values returns EA, unpack for backward-compat if getattr(self.dtype, "tz", None) is None: result = np.asarray(result)

…e.unique

Dr-Irv · 2021-03-16T12:03:31Z

@simonjayhawkins I believe I've addressed all your requests on this. Anything else to do?

simonjayhawkins · 2021-03-16T13:21:12Z

Thanks @Dr-Irv for addressing the comments. I want to make another thorough pass before approving.

jreback · 2021-03-16T14:32:02Z

pandas/core/arrays/base.py

+        # error: Argument "dtype" to "array" has incompatible type
+        # "Union[ExtensionDtype, dtype[Any]]"; expected "Union[dtype[Any], None,
+        # type, _SupportsDType, str, Union[Tuple[Any, int], Tuple[Any,
+        # Union[int, Sequence[int]]], List[Any], _DTypeDict, Tuple[Any, Any]]]"


can't you just assert that this is a NpDtype here?

We decided to address astype() (where the above code is) in a later PR. That will remove the above.

see previous comment, remove the added type and this ignore will not be needed yet.

Dr-Irv · 2021-04-03T15:48:09Z

@simonjayhawkins pinging you again on review of this. Thanks!

simonjayhawkins · 2021-04-27T15:52:49Z

@Dr-Irv can you resolve conflicts.

Dr-Irv · 2021-04-27T20:47:02Z

@Dr-Irv can you resolve conflicts.

@simonjayhawkins I've done this. Note that with respect to what was in #41097, I changed the name of the TypeVar to be ExtensionDtypeT rather than E to follow convention that we use elsewhere.

Note that the way that I use pyright didn't pick that up because the example was from CategoricalDtype and I was only verifying the base extension types.

simonjayhawkins

Thanks @Dr-Irv

simonjayhawkins · 2021-04-28T13:49:27Z

pandas/core/array_algos/putmask.py

@@ -191,7 +191,7 @@ def extract_bool_array(mask: ArrayLike) -> np.ndarray:
        # We could have BooleanArray, Sparse[bool], ...
        #  Except for BooleanArray, this is equivalent to just
        #  np.asarray(mask, dtype=bool)
-        mask = mask.to_numpy(dtype=bool, na_value=False)
+        mask = mask.to_numpy(dtype=np.dtype(bool), na_value=False)


This shouldn't need to be changed. i've opened #41185 as a precursor to fix these false positives.

simonjayhawkins · 2021-04-28T14:01:06Z

pandas/core/arrays/base.py

        copy: bool = False,
-        na_value=lib.no_default,
+        na_value: Any | None = lib.no_default,


None seems redundant here.

simonjayhawkins · 2021-04-28T14:11:38Z

pandas/core/arrays/base.py

@@ -510,8 +508,7 @@ def nbytes(self) -> int:
    # ------------------------------------------------------------------------
    # Additional Methods
    # ------------------------------------------------------------------------
-
-    def astype(self, dtype, copy=True):
+    def astype(self, dtype: Dtype, copy: bool = True):


when adding types, we should ensure that the types match the docstrings. From the PR title and previous discussion, IIUC astype was being done separately.

I would revert this for now until this method is sorted properly.

see also #41018 (comment) and response.

The docstring states we return np.ndarray and the one-liner suggests that too. We sometimes also return an ExtensionArray, this is dependent on the type of dtype.

my concern is that if we add the type now, this may get forgotten.

simonjayhawkins · 2021-04-28T14:16:18Z

pandas/core/arrays/base.py

+        # error: Argument "dtype" to "array" has incompatible type
+        # "Union[ExtensionDtype, dtype[Any]]"; expected "Union[dtype[Any], None,
+        # type, _SupportsDType, str, Union[Tuple[Any, int], Tuple[Any,
+        # Union[int, Sequence[int]]], List[Any], _DTypeDict, Tuple[Any, Any]]]"


see previous comment, remove the added type and this ignore will not be needed yet.

simonjayhawkins · 2021-04-28T14:24:57Z

pandas/core/arrays/base.py

+        *args: Any,
+        **kwargs: Any,


I would leave these untyped for now.

These are passed through to numpy.argsort, so we should be able to be more precise.

try to avoid adding Any where we maybe able to more precise with further investigation.

side note: these actually don't seem to be passed in the base implementation.

simonjayhawkins · 2021-04-28T15:59:00Z

pandas/core/arrays/string_arrow.py

@@ -696,7 +693,7 @@ def value_counts(self, dropna: bool = True) -> Series:

    _str_na_value = ArrowStringDtype.na_value

-    def _str_map(self, f, na_value=None, dtype: Dtype | None = None):
+    def _str_map(self, f, na_value=None, dtype: DtypeArg | None = None):


pandas/core/dtypes/base.py

simonjayhawkins · 2021-04-28T17:24:53Z

pandas/core/indexes/base.py

+        # type "Union[ExtensionArray, ndarray]"; expected "Sequence[Any]"
+        locs = self._values[mask].searchsorted(
+            where._values, side="right"  # type: ignore[arg-type]
+        )


see comments regarding using Sequence for ArrayLike

Dr-Irv · 2021-05-02T00:11:14Z

Closing this as I've started making even smaller PRs and tracking status here: #39501 (comment) in the comment from the original big PR

Dr-Irv and others added 29 commits January 23, 2021 11:33

small typing fixes

f2c52a4

fix ExtensionArray and EXtensionDtype

d7ff8d3

merge with master

49fa06e

fixes for delete, isin, unique

03b2c4a

fix import of Literal

3e19958

remove quotes on ExtensionDType.construct_from_string

6861901

move numpy workaround to _typing.py

9be6486

remove numpy dummy

260b367

remove extra line in _typing

6276725

Merge remote-tracking branch 'upstream/master' into extensiontyping

4dafaca

import Literal

8b2cee2

Merge remote-tracking branch 'upstream/master' into extensiontyping

3a7d839

merge with master

a21bb60

isort precommit fix

8cd6b76

fix interval.repeat() typing

e0e0131

overload for __getitem__ and use pattern with ExtensionArrayT as self…

6a6a21f

… and return type

lose less ExtensionArrayT. Make registry private. consolidate overload

bf753e6

remove ExtensionArray typing of self

c9795a5

Merge remote-tracking branch 'upstream/master' into extensiontyping

d452842

merge with upstream/master

3c2c78b

make extension arrays work with new typing, fixing astype and to_numpy

548c198

fix Literal import

db8ed9b

fix logic in ensure_int_or_float

f8191f8

fix conflict with master

575645f

fix typing in groupby to_numpy call

6f8fcb5

fix groupby again. Allow kwargs for extension to_numpy

3ea2420

Merge remote-tracking branch 'upstream/master' into extensiontyping

c83a628

fixes for merge with master

5bb24d4

remove astype and isin changes

ad1ab3b

Dr-Irv requested a review from simonjayhawkins March 13, 2021 17:27

simonjayhawkins reviewed Mar 13, 2021

View reviewed changes

pandas/core/arrays/base.py Show resolved Hide resolved

don't cast in astype. TODO for overload of astype

1196132

simonjayhawkins reviewed Mar 13, 2021

View reviewed changes

Dr-Irv added 3 commits March 13, 2021 17:44

remove private registry, getitem overloads, typevar on DateTimeScalar

66d5da4

Remove List[Any] from getitem

5411998

remove spacing change in _mixins.py and __getitem__

9b7481d

jbrockmendel reviewed Mar 13, 2021

View reviewed changes

pandas/io/formats/format.py Outdated Show resolved Hide resolved

jbrockmendel reviewed Mar 13, 2021

View reviewed changes

remove cast in io/formats. Change isinstance check in pandas.core.bas…

4bd3422

…e.unique

simonjayhawkins added the Typing type annotations, mypy/pyright type checking label Mar 14, 2021

Dr-Irv mentioned this pull request Mar 14, 2021

TYP: Update type hints for ExtensionArray and ExtensionDtype #39501

Closed

merge with master to resolve conflicts

3b1ff79

jreback reviewed Mar 16, 2021

View reviewed changes

Dr-Irv added 2 commits April 2, 2021 23:38

merge with master, remove more ignores

5719daa

remove mypy comments from format

dbfb3a2

Dr-Irv mentioned this pull request Apr 24, 2021

TYP: Fix typehints for ExtensionDtype #41097

Merged

resolve conflicts with master

d41bf91

Dr-Irv mentioned this pull request Apr 28, 2021

TYP: use DTypeLike alias from numpy.typing #41185

Closed

simonjayhawkins reviewed Apr 28, 2021

View reviewed changes

Dr-Irv marked this pull request as draft April 28, 2021 18:41

Dr-Irv mentioned this pull request Apr 28, 2021

TYP: Fix typing in ExtensionDtype registry #41203

Merged

Dr-Irv closed this May 2, 2021

Dr-Irv deleted the limitextensiontyping branch February 13, 2023 20:50

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix typing for extension arrays and extension dtypes without isin and astype #40421

Fix typing for extension arrays and extension dtypes without isin and astype #40421

Dr-Irv commented Mar 13, 2021

Dr-Irv commented Mar 13, 2021 •

edited

simonjayhawkins Mar 13, 2021

Dr-Irv Mar 13, 2021

simonjayhawkins Mar 13, 2021

Dr-Irv Mar 13, 2021

simonjayhawkins Mar 13, 2021

Dr-Irv Mar 13, 2021

jbrockmendel Mar 13, 2021

Dr-Irv Mar 14, 2021

jbrockmendel Mar 13, 2021

Dr-Irv Mar 14, 2021

Dr-Irv commented Mar 16, 2021

simonjayhawkins commented Mar 16, 2021

jreback Mar 16, 2021

Dr-Irv Mar 16, 2021

simonjayhawkins Apr 28, 2021

Dr-Irv commented Apr 3, 2021

simonjayhawkins commented Apr 27, 2021

Dr-Irv commented Apr 27, 2021 •

edited

simonjayhawkins left a comment

simonjayhawkins Apr 28, 2021

simonjayhawkins Apr 28, 2021

simonjayhawkins Apr 28, 2021

simonjayhawkins Apr 28, 2021

simonjayhawkins Apr 28, 2021

simonjayhawkins Apr 28, 2021

simonjayhawkins Apr 28, 2021

Dr-Irv commented May 2, 2021

Fix typing for extension arrays and extension dtypes without isin and astype #40421

Fix typing for extension arrays and extension dtypes without isin and astype #40421

Conversation

Dr-Irv commented Mar 13, 2021

Dr-Irv commented Mar 13, 2021 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Dr-Irv commented Mar 16, 2021

simonjayhawkins commented Mar 16, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Dr-Irv commented Apr 3, 2021

simonjayhawkins commented Apr 27, 2021

Dr-Irv commented Apr 27, 2021 • edited

simonjayhawkins left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Dr-Irv commented May 2, 2021

Dr-Irv commented Mar 13, 2021 •

edited

Dr-Irv commented Apr 27, 2021 •

edited