Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add typing annotation to assert_index_equal #26535

Merged
merged 11 commits into from
Jun 8, 2019
15 changes: 12 additions & 3 deletions pandas/util/testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import string
import tempfile
import traceback
from typing import Union, cast
import warnings
import zipfile

Expand Down Expand Up @@ -515,9 +516,14 @@ def equalContents(arr1, arr2):
return frozenset(arr1) == frozenset(arr2)


def assert_index_equal(left, right, exact='equiv', check_names=True,
check_less_precise=False, check_exact=True,
check_categorical=True, obj='Index'):
def assert_index_equal(left: Index,
right: Index,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of a raw Index, is it possible to add a typing.TypeVar to _typing.py:

AnyIndex = TypeVar('AnyIndex ', pd.Index, pd.RangeIndex, pd.Int64Index, pd.MultiIndex, + other index types...)

and then use that here (I'm not super aqianted with typing, so could be wrong)?

That would avoid the need to adapt the code below.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is ok as all of these are Index subclasses

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The point was that AnyIndex would capture .levels and other attributes that are not set on the base Index.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is OK too. MyPy would still complain about those attributes not existing on Index

exact: Union[bool, str] = 'equiv',
check_names: bool = True,
check_less_precise: Union[bool, int] = False,
check_exact: bool = True,
check_categorical: bool = True,
obj: str = 'Index') -> None:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a general comment but NoReturn may be applicable here at some point when we drop 3.5, though not sure if Optional[NoReturn] makes sense

"""Check that left and right Index are equal.

Parameters
Expand Down Expand Up @@ -588,6 +594,9 @@ def _get_ilevel_values(index, level):

# MultiIndex special comparison for little-friendly error messages
if left.nlevels > 1:
left = cast(MultiIndex, left)
right = cast(MultiIndex, right)

for level in range(left.nlevels):
# cannot use get_level_values here because it can change dtype
llevel = _get_ilevel_values(left, level)
Expand Down