issue23: Fix __getitem__ and __setitem__ signatures (#24)

* Fixed issue with __getitem__ and __setitem__ signatures.

* Test fix

Author: zkrolikowski-vl

setup.py

@@ -21,7 +21,7 @@ def list_packages(source_path: str = src_path) -> None:
 setup(
     name="pandas-stubs",
     package_dir={"": src_path},
-    version="1.0.4.5",
+    version="1.0.4.6",
     description="Type annotations for Pandas",
     long_description=(open("README.md").read()
                       if os.path.exists("README.md") else ""),

tests/snippets/test_frame.py

@@ -1,4 +1,8 @@
+# flake8: noqa: F841
+
 import tempfile
+from pathlib import Path
+
 import pandas as pd
 import numpy as np
 
@@ -12,20 +16,46 @@ def test_types_init() -> None:
 
 def test_types_csv() -> None:
     df = pd.DataFrame(data={'col1': [1, 2], 'col2': [3, 4]})
-    with tempfile.TemporaryFile() as file:
-        df.to_csv()
-        df.to_csv(file.name)
-        df.to_csv(file)
-        pd.read_csv()
-        pd.read_csv(file)
-        pd.read_csv(file.name)
+    csv_df: str = df.to_csv()
 
-
-def test_types_select() -> None:
-    df = pd.DataFrame(data={'col1': [1, 2], 'col2': [3, 4]})
+    with tempfile.NamedTemporaryFile() as file:
+        df.to_csv(file.name)
+        df2: pd.DataFrame = pd.read_csv(file.name)
+
+    with tempfile.NamedTemporaryFile() as file:
+        df.to_csv(Path(file.name))
+        df3: pd.DataFrame = pd.read_csv(Path(file.name))
+
+def test_types_getitem() -> None:
+    df = pd.DataFrame(data={'col1': [1, 2], 'col2': [3, 4], 5: [6, 7]})
+    i = pd.Index(['col1', 'col2'])
+    s = pd.Series(['col1', 'col2'])
+    select_df = pd.DataFrame({'col1': [True, True], 'col2': [False, True]})
+    a = np.array(['col1', 'col2'])
     df['col1']
-    df[0]
+    df[5]
+    df[['col1', 'col2']]
     df[1:]
+    df[s]
+    df[a]
+    df[select_df]
+    df[i]
+
+
+def test_types_setitem() -> None:
+    df = pd.DataFrame(data={'col1': [1, 2], 'col2': [3, 4], 5: [6, 7]})
+    i = pd.Index(['col1', 'col2'])
+    s = pd.Series(['col1', 'col2'])
+    select_df = pd.DataFrame({'col1': [True, True], 'col2': [False, True]})
+    a = np.array(['col1', 'col2'])
+    df['col1'] = [1, 2]
+    df[5] = [5, 6]
+    df[['col1', 'col2']] = [[1, 2], [3, 4]]
+    df[1:] = {'col1': {1: 'a'}, 'col2': {1: 'b'}, 5: {1: 'c'}}
+    df[s] = [5, 6]
+    df[a] = [[1, 2], [3, 4]]
+    df[select_df] = [1, 2, 3]
+    df[i] = [8, 9]
 
 
 def test_types_iloc_iat() -> None:

tests/snippets/test_series.py

@@ -1,4 +1,7 @@
+# flake8: noqa: F841
+
 import tempfile
+from pathlib import Path
 
 import pandas as pd
 import numpy as np
@@ -16,18 +19,18 @@ def test_types_init() -> None:
 
 def test_types_csv() -> None:
     s = pd.Series(data=[1, 2, 3])
-    with tempfile.TemporaryFile() as file:
-        s.to_csv()
+    csv_df: str = s.to_csv()
+
+    with tempfile.NamedTemporaryFile() as file:
         s.to_csv(file.name)
-        s.to_csv(file)
-        pd.read_csv()
-        pd.read_csv(file)
-        pd.read_csv(file.name)
+        s2: pd.Series = pd.read_csv(file.name)
 
+    with tempfile.NamedTemporaryFile() as file:
+        s.to_csv(Path(file.name))
+        s3: pd.Series = pd.read_csv(Path(file.name))
 
 def test_types_select() -> None:
     s = pd.Series(data={'row1': 1, 'row2': 2})
-    s['col1']
     s[0]
     s[1:]
 
@@ -53,7 +56,7 @@ def test_types_loc_at() -> None:
 def test_types_boolean_indexing() -> None:
     s = pd.Series([0, 1, 2])
     s[s > 1]
-    s[pd]
+    s[s]
 
 
 def test_types_head_tail() -> None:
@@ -230,7 +233,7 @@ def test_types_var() -> None:
 def test_types_std() -> None:
     s = pd.Series([-10, 2, 3, 10])
     s.std()
-    s.std(axis=1, ddof=1)
+    s.std(axis=0, ddof=1)
     s.std(skipna=True, numeric_only=False)
 
 

third_party/3/pandas/core/frame.pyi

@@ -73,20 +73,15 @@ class DataFrame(NDFrame):
     def __sub__(self, other: Union[Scalar, DataFrame]) -> DataFrame: ...
     def __matmul__(self, other: Union[FrameOrSeries, ArrayLike]) -> FrameOrSeries: ...
     def __rmatmul__(self, other: Union[FrameOrSeries, ArrayLike]) -> FrameOrSeries: ...
+    # Type ignoring the first declaration because for some reason mypy thinks that np.ndarray overlaps Column
     @overload
-    def __getitem__(self, key: Column) -> Series: ...
+    def __getitem__(self, key: Column) -> Series: ...  # type: ignore
     @overload
-    def __getitem__(self, key: Sequence[Column]) -> Series: ...
+    def __getitem__(self, key: Union[Series, Index, DataFrame, List[Column], slice, np.ndarray]) -> DataFrame: ...
     @overload
-    def __getitem__(self, key: DataFrame) -> DataFrame: ...
+    def __setitem__(self, key: Column, value: Any) -> DataFrame: ...
     @overload
-    def __getitem__(self, key: slice) -> DataFrame: ...
-    @overload
-    def __setitem__(self, key: Column, value: Any) -> None: ...
-    @overload
-    def __setitem__(self, key: Sequence[Column], value: Any) -> None: ...
-    @overload
-    def __setitem__(self, key: slice, value: Any) -> None: ...
+    def __setitem__(self, key: Union[Series, Index, DataFrame, List[Column], slice, np.ndarray], value: Any) -> DataFrame: ...
     @property
     def axes(self) -> List[Index]: ...
     @property