Fix concat, add tests (#84)

Author: zkrolikowski-vl

tests/snippets/test_pandas.py

@@ -24,6 +24,14 @@ def test_types_concat() -> None:
     pd.concat([s, s2], keys=['first', 'second'], sort=True)
     pd.concat([s, s2], keys=['first', 'second'], names=["source", "row"])
 
+    # Depends on the axis
+    rs1: Union[pd.Series, pd.DataFrame] = pd.concat({'a': s, 'b': s2})
+    rs1a: Union[pd.Series, pd.DataFrame] = pd.concat({'a': s, 'b': s2}, axis=1)
+    rs2: Union[pd.Series, pd.DataFrame] = pd.concat({1: s, 2: s2})
+    rs2a: Union[pd.Series, pd.DataFrame] = pd.concat({1: s, 2: s2}, axis=1)
+    rs3: Union[pd.Series, pd.DataFrame] = pd.concat({1: s, None: s2})
+    rs3a: Union[pd.Series, pd.DataFrame] = pd.concat({1: s, None: s2}, axis=1)
+
     df = pd.DataFrame(data={'col1': [1, 2], 'col2': [3, 4]})
     df2 = pd.DataFrame(data={'col1': [10, 20], 'col2': [30, 40]})
 
@@ -34,3 +42,7 @@ def test_types_concat() -> None:
 
     result: pd.DataFrame = pd.concat({"a": pd.DataFrame([1, 2, 3]), "b": pd.DataFrame([4, 5, 6])}, axis=1)
     result2: Union[pd.DataFrame, pd.Series] = pd.concat({"a": pd.Series([1, 2, 3]), "b": pd.Series([4, 5, 6])}, axis=1)
+
+    rdf1: pd.DataFrame = pd.concat({'a': df, 'b': df2})
+    rdf2: pd.DataFrame = pd.concat({1: df, 2: df2})
+    rdf3: pd.DataFrame = pd.concat({1: df, None: df2})

third_party/3/pandas/core/reshape/concat.pyi

@@ -4,17 +4,26 @@ from pandas.core.arrays.categorical import factorize_from_iterable as factorize_
 from pandas.core.generic import NDFrame as NDFrame
 from pandas.core.indexes.api import all_indexes_same as all_indexes_same, ensure_index as ensure_index, get_consensus_names as get_consensus_names, get_objs_combined_axis as get_objs_combined_axis
 from pandas.core.internals import concatenate_block_managers as concatenate_block_managers
-from typing import Any, Iterable, Mapping, overload
+from typing import Any, Iterable, Mapping, overload, Optional
+
 
 # For some reason mypy won't read this declaration correctly if we use 2 overloads with union
 @overload
 def concat(objs: Iterable[DataFrame], axis: Axis = ..., join: str=..., ignore_index: bool=..., keys: Any = ..., levels: Any = ..., names: Any = ..., verify_integrity: bool=..., sort: bool=..., copy: bool=...) -> DataFrame: ...
 @overload
+def concat(objs: Mapping[Optional[int], DataFrame], axis: Axis = ..., join: str=..., ignore_index: bool=..., keys: Any = ..., levels: Any = ..., names: Any = ..., verify_integrity: bool=..., sort: bool=..., copy: bool=...) -> DataFrame: ...
+@overload
+def concat(objs: Mapping[Optional[str], DataFrame], axis: Axis = ..., join: str=..., ignore_index: bool=..., keys: Any = ..., levels: Any = ..., names: Any = ..., verify_integrity: bool=..., sort: bool=..., copy: bool=...) -> DataFrame: ...
+@overload
 def concat(objs: Mapping[Label, DataFrame], axis: Axis = ..., join: str=..., ignore_index: bool=..., keys: Any = ..., levels: Any = ..., names: Any = ..., verify_integrity: bool=..., sort: bool=..., copy: bool=...) -> DataFrame: ...
 @overload
 def concat(objs: Iterable[Series], axis: Axis = ..., join: str=..., ignore_index: bool=..., keys: Any = ..., levels: Any = ..., names: Any = ..., verify_integrity: bool=..., sort: bool=..., copy: bool=...) -> FrameOrSeries: ...
 @overload
 def concat(objs: Mapping[Label, Series], axis: Axis = ..., join: str=..., ignore_index: bool=..., keys: Any = ..., levels: Any = ..., names: Any = ..., verify_integrity: bool=..., sort: bool=..., copy: bool=...) -> FrameOrSeries: ...
+@overload
+def concat(objs: Mapping[Optional[int], Series], axis: Axis = ..., join: str=..., ignore_index: bool=..., keys: Any = ..., levels: Any = ..., names: Any = ..., verify_integrity: bool=..., sort: bool=..., copy: bool=...) -> FrameOrSeries: ...
+@overload
+def concat(objs: Mapping[Optional[str], Series], axis: Axis = ..., join: str=..., ignore_index: bool=..., keys: Any = ..., levels: Any = ..., names: Any = ..., verify_integrity: bool=..., sort: bool=..., copy: bool=...) -> FrameOrSeries: ...
 
 class _Concatenator:
     intersect: bool = ...