TYP: Add annotation for df.pivot
State: closed

- [ ] xref #30928 
- [ ] tests added / passed
- [ ] passes `black pandas`
- [ ] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [ ] whatsnew entry


Labels:
name: Typing
description: type annotations, mypy/pyright type checking

Comments:
Hello @charlesdong1991! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:





There are currently no PEP 8 issues detected in this Pull Request. Cheers! :beers: 

##### Comment last updated at 2020-05-18 19:53:56 UTC
@charlesdong1991 can you fix up conflict and address comments?
Hi, @WillAyd i have tried quite many times, but could not find a working solution than using `Axes`, if using `Label` as suggested, I will see errors below:
```
pandas/core/reshape/pivot.py:442: error: No overload variant of "list" matches argument type "Hashable"
pandas/core/reshape/pivot.py:442: note: Possible overload variant:
pandas/core/reshape/pivot.py:442: note:     def [_T] list(self, iterable: Iterable[_T]) -> List[_T]
pandas/core/reshape/pivot.py:442: note:     <1 more non-matching overload not shown>
pandas/core/reshape/pivot.py:445: error: Argument 1 to "extend" of "list" has incompatible type "Hashable"; expected "Iterable[Any]"
pandas/core/reshape/pivot.py:453: error: "Hashable" has no attribute "__iter__" (not iterable)
pandas/core/reshape/pivot.py:457: error: "Hashable" has no attribute "__iter__" (not iterable)
pandas/core/reshape/pivot.py:464: error: Argument "columns" to "DataFrame" has incompatible type "Hashable"; expected "Optional[Collection[Any]]"
```

do you have any idea of how to work around it? some hints would be super helpful and appreciated! @WillAyd thanks a lot!

@WillAyd @simonjayhawkins @jbrockmendel 

thanks for all the comments, and I have adjusted a bit based on my understanding of your preference: to use `# type ignore` for now and comment to change once mypy gets better.

Feel free to take another look and any comments are very welcomed!
Given the troubles here and the amount of "type: ignore" going in, im not sure this is worthwhile ATM.
> not a fan of all of the type ignores

we have the choice of ignores, ignores with error codes, casts and asserts.

my preference here would be casts (only after is_list_like calls) xref #32785

```diff
diff --git a/pandas/core/reshape/pivot.py b/pandas/core/reshape/pivot.py
index e49dd4c9b..855189c43 100644
--- a/pandas/core/reshape/pivot.py
+++ b/pandas/core/reshape/pivot.py
@@ -1,4 +1,14 @@
-from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    cast,
+)
 
 import numpy as np
 
@@ -427,24 +437,28 @@ def _convert_by(by):
 @Appender(_shared_docs["pivot"], indents=1)
 def pivot(
     data: "DataFrame",
-    index: Optional[Union[Label, Sequence[Optional[Label]]]] = None,
-    columns: Optional[Union[Label, Sequence[Optional[Label]]]] = None,
-    values: Optional[Union[Label, Sequence[Optional[Label]]]] = None,
+    index: Optional[Union[Label, Sequence[Label]]] = None,
+    columns: Optional[Union[Label, Sequence[Label]]] = None,
+    values: Optional[Union[Label, Sequence[Label]]] = None,
 ) -> "DataFrame":
     if columns is None:
         raise TypeError("pivot() missing 1 required argument: 'columns'")
-    columns = columns if is_list_like(columns) else [columns]
+    if is_list_like(columns):
+        columns = cast(Sequence[Label], columns)
+        columns = columns
+    else:
+        columns = [columns]
 
     if values is None:
-        cols: List[Optional[Sequence]] = []
+        cols: List[Label] = []
         if index is None:
             pass
         elif is_list_like(index):
-            # Remove type ignore once mypy-5206 is implemented, same for below
-            cols = list(index)  # type: ignore
+            index = cast(Sequence[Label], index)
+            cols = list(index)
         else:
-            cols = [index]  # type: ignore
-        cols.extend(columns)  # type: ignore
+            cols = [index]
+        cols.extend(columns)
 
         append = index is None
         indexed = data.set_index(cols, append=append)
@@ -452,18 +466,20 @@ def pivot(
         if index is None:
             idx_list = [Series(data.index, name=data.index.name)]
         elif is_list_like(index):
-            idx_list = [data[idx] for idx in index]  # type: ignore
+            index = cast(Sequence[Label], index)
+            idx_list = [data[idx] for idx in index]
         else:
             idx_list = [data[index]]
 
-        data_columns = [data[col] for col in columns]  # type: ignore
+        data_columns = [data[col] for col in columns]
         idx_list.extend(data_columns)
         mi_index = MultiIndex.from_arrays(idx_list)
 
         if is_list_like(values) and not isinstance(values, tuple):
             # Exclude tuple because it is seen as a single column name
+            values = cast(Sequence[Label], values)
             indexed = data._constructor(
-                data[values]._values, index=mi_index, columns=values  # type: ignore
+                data[values]._values, index=mi_index, columns=values
             )
         else:
             indexed = data._constructor_sliced(data[values]._values, index=mi_index)
```

it looks like the use of type: ignores here masked an incorrect type definition of `cols`

each cast used here could instead be replaced with a runtime assert until resolution of #32785

This diff uses 4 casts corresponding to 4 is_list_like calls compared with 6 ignores for which the reasons are not immediately obvious.
@charlesdong1991 whats the status here?
I have updated and it seems `cast` indeed work here. 

Thanks very much! @simonjayhawkins 
@jreback i use `convert_to_list_like` for the first one, keep your second comment as is, and pls let me know if you think it's okay
many thanks for the help on annotations @simonjayhawkins i think it's good to go now!
Thanks @charlesdong1991 .lgtm. I _think_ if you type the return type of MultiIndex.from_arrays you could revert the variable name changes idx_list and mi_index and reduce the diff.

```
    def from_arrays(cls, arrays, sortorder=None, names=lib.no_default) -> "MultiIndex":
```
ahh, yeah, indeed seems pass mypy locally. 

very nice!! @simonjayhawkins thanks!
thanks @charlesdong1991