From 9fb27cdc1987f9f46ce90b6adcdf90ef3c676dce Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Fri, 5 May 2023 22:21:14 +0530 Subject: [PATCH 1/7] added `mode` keyword to `DataFrame.to_json` --- pandas-stubs/_typing.pyi | 4 +--- pandas-stubs/core/arrays/base.pyi | 5 +++- pandas-stubs/core/frame.pyi | 40 +++++++++++++++++++++++++++++-- tests/test_frame.py | 14 +++++++++++ 4 files changed, 57 insertions(+), 6 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 6bdb6f91e..1d33a58be 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -445,9 +445,7 @@ SortKind: TypeAlias = Literal["quicksort", "mergesort", "heapsort", "stable"] NaPosition: TypeAlias = Literal["first", "last"] JoinHow: TypeAlias = Literal["left", "right", "outer", "inner"] MergeHow: TypeAlias = JoinHow | Literal["cross"] -JsonFrameOrient: TypeAlias = Literal[ - "split", "records", "index", "columns", "values", "table" -] +JsonFrameOrient: TypeAlias = Literal["split", "index", "columns", "values", "table"] JsonSeriesOrient: TypeAlias = Literal["split", "records", "index", "table"] TimestampConvention: TypeAlias = Literal["start", "end", "s", "e"] diff --git a/pandas-stubs/core/arrays/base.pyi b/pandas-stubs/core/arrays/base.pyi index 7141f70ff..d81272928 100644 --- a/pandas-stubs/core/arrays/base.pyi +++ b/pandas-stubs/core/arrays/base.pyi @@ -1,4 +1,7 @@ -from typing import Any +from typing import ( + Any, + Literal, +) import numpy as np from typing_extensions import Self diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index b94912442..df08158c2 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -2019,10 +2019,11 @@ class DataFrame(NDFrame, OpsMixin): date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: _bool = ..., + lines: Literal["False"] = ..., compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., + mode: Literal["w"] = ..., ) -> None: ... @overload def to_json( @@ -2035,10 +2036,45 @@ class DataFrame(NDFrame, OpsMixin): date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: _bool = ..., + lines: Literal["False"] = ..., compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., + mode: Literal["w"] = ..., + ) -> _str: ... + @overload + def to_json( + self, + path_or_buf: FilePath | WriteBuffer[str], + orient: Literal["records"] = ..., + date_format: Literal["epoch", "iso"] | None = ..., + double_precision: int = ..., + force_ascii: _bool = ..., + date_unit: Literal["s", "ms", "us", "ns"] = ..., + default_handler: Callable[[Any], _str | float | _bool | list | dict] + | None = ..., + lines: Literal["True"] = ..., + compression: CompressionOptions = ..., + index: _bool = ..., + indent: int | None = ..., + mode: Literal["a"] = ..., + ) -> None: ... + @overload + def to_json( + self, + path_or_buf: None = ..., + orient: Literal["records"] = ..., + date_format: Literal["epoch", "iso"] | None = ..., + double_precision: int = ..., + force_ascii: _bool = ..., + date_unit: Literal["s", "ms", "us", "ns"] = ..., + default_handler: Callable[[Any], _str | float | _bool | list | dict] + | None = ..., + lines: Literal["True"] = ..., + compression: CompressionOptions = ..., + index: _bool = ..., + indent: int | None = ..., + mode: Literal["a"] = ..., ) -> _str: ... @overload def to_string( diff --git a/tests/test_frame.py b/tests/test_frame.py index 94f279563..1878ed0cf 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2578,3 +2578,17 @@ def test_convert_dtypes_dtype_backend() -> None: df = pd.DataFrame({"A": [1, 2, 3, 4], "B": [3, 4, 5, 6]}) dfn = df.convert_dtypes(dtype_backend="numpy_nullable") check(assert_type(dfn, pd.DataFrame), pd.DataFrame) + + +def test_to_json_mode() -> None: + df = pd.DataFrame( + [["a", "b"], ["c", "d"]], + index=["row 1", "row 2"], + columns=["col 1", "col 2"], + ) + result = df.to_json(orient="records", lines="True", mode="a") + result1 = df.to_json(orient="split", mode="w") + result2 = df.to_json(orient="columns", mode="w") + check(assert_type(result, str), str) + check(assert_type(result1, str), str) + check(assert_type(result2, str), str) From 1146dcb10b98a2515e8503e8e8b054bf75f38655 Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Fri, 5 May 2023 22:28:05 +0530 Subject: [PATCH 2/7] Update base.pyi --- pandas-stubs/core/arrays/base.pyi | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pandas-stubs/core/arrays/base.pyi b/pandas-stubs/core/arrays/base.pyi index d81272928..7141f70ff 100644 --- a/pandas-stubs/core/arrays/base.pyi +++ b/pandas-stubs/core/arrays/base.pyi @@ -1,7 +1,4 @@ -from typing import ( - Any, - Literal, -) +from typing import Any import numpy as np from typing_extensions import Self From addd3a07e8a98380d63bad0048568acd8a5f0b8c Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Wed, 10 May 2023 14:26:39 +0530 Subject: [PATCH 3/7] added overloads and tests --- pandas-stubs/_typing.pyi | 4 +++- pandas-stubs/core/frame.pyi | 34 +++++++++++++++++--------------- pandas-stubs/core/series.pyi | 38 ++++++++++++++++++++++++++++++++++++ tests/test_frame.py | 7 ++++++- tests/test_series.py | 15 ++++++++++++++ 5 files changed, 80 insertions(+), 18 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 1d33a58be..6bdb6f91e 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -445,7 +445,9 @@ SortKind: TypeAlias = Literal["quicksort", "mergesort", "heapsort", "stable"] NaPosition: TypeAlias = Literal["first", "last"] JoinHow: TypeAlias = Literal["left", "right", "outer", "inner"] MergeHow: TypeAlias = JoinHow | Literal["cross"] -JsonFrameOrient: TypeAlias = Literal["split", "index", "columns", "values", "table"] +JsonFrameOrient: TypeAlias = Literal[ + "split", "records", "index", "columns", "values", "table" +] JsonSeriesOrient: TypeAlias = Literal["split", "records", "index", "table"] TimestampConvention: TypeAlias = Literal["start", "end", "s", "e"] diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index df08158c2..7fb7e6b3d 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -2012,70 +2012,72 @@ class DataFrame(NDFrame, OpsMixin): def to_json( self, path_or_buf: FilePath | WriteBuffer[str], - orient: JsonFrameOrient | None = ..., + *, + orient: Literal["records"], date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., force_ascii: _bool = ..., date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: Literal["False"] = ..., + lines: Literal[True], compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., - mode: Literal["w"] = ..., + mode: Literal["a"], ) -> None: ... @overload def to_json( self, path_or_buf: None = ..., - orient: JsonFrameOrient | None = ..., + *, + orient: Literal["records"], date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., force_ascii: _bool = ..., date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: Literal["False"] = ..., + lines: Literal[True], compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., - mode: Literal["w"] = ..., + mode: Literal["a"], ) -> _str: ... @overload def to_json( self, - path_or_buf: FilePath | WriteBuffer[str], - orient: Literal["records"] = ..., + path_or_buf: None = ..., + orient: JsonFrameOrient | None = ..., date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., force_ascii: _bool = ..., date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: Literal["True"] = ..., + lines: _bool = ..., compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., - mode: Literal["a"] = ..., - ) -> None: ... + mode: Literal["w"] = ..., + ) -> _str: ... @overload def to_json( self, - path_or_buf: None = ..., - orient: Literal["records"] = ..., + path_or_buf: FilePath | WriteBuffer[str], + orient: JsonFrameOrient | None = ..., date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., force_ascii: _bool = ..., date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: Literal["True"] = ..., + lines: _bool = ..., compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., - mode: Literal["a"] = ..., - ) -> _str: ... + mode: Literal["w"] = ..., + ) -> None: ... @overload def to_string( self, diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 41c1e9f27..62feaf9d1 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -500,6 +500,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., + mode: Literal["w"] = ..., ) -> None: ... @overload def to_json( @@ -516,6 +517,43 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., + mode: Literal["w"] = ..., + ) -> _str: ... + @overload + def to_json( + self, + *, + path_or_buf: FilePath | WriteBuffer[str], + orient: Literal["records"], + date_format: Literal["epoch", "iso"] | None = ..., + double_precision: int = ..., + force_ascii: _bool = ..., + date_unit: Literal["s", "ms", "us", "ns"] = ..., + default_handler: Callable[[Any], _str | float | _bool | list | dict] + | None = ..., + lines: Literal[True], + compression: CompressionOptions = ..., + index: _bool = ..., + indent: int | None = ..., + mode: Literal["a"], + ) -> None: ... + @overload + def to_json( + self, + *, + path_or_buf: None = ..., + orient: Literal["records"], + date_format: Literal["epoch", "iso"] | None = ..., + double_precision: int = ..., + force_ascii: _bool = ..., + date_unit: Literal["s", "ms", "us", "ns"] = ..., + default_handler: Callable[[Any], _str | float | _bool | list | dict] + | None = ..., + lines: Literal[True], + compression: CompressionOptions = ..., + index: _bool = ..., + indent: int | None = ..., + mode: Literal["a"], ) -> _str: ... def to_xarray(self) -> xr.DataArray: ... def items(self) -> Iterable[tuple[Hashable, S1]]: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 1878ed0cf..4a8eccf2a 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2586,9 +2586,14 @@ def test_to_json_mode() -> None: index=["row 1", "row 2"], columns=["col 1", "col 2"], ) - result = df.to_json(orient="records", lines="True", mode="a") + result = df.to_json(orient="records", lines=True, mode="a") result1 = df.to_json(orient="split", mode="w") result2 = df.to_json(orient="columns", mode="w") check(assert_type(result, str), str) check(assert_type(result1, str), str) check(assert_type(result2, str), str) + if TYPE_CHECKING_INVALID_USAGE: + result3 = df.to_json(orient="records", lines=False, mode="a") # type: ignore[call-overload] + result4 = df.to_json(orient="records", mode="w") + check(assert_type(result3, Any), str) + check(assert_type(result4, str), str) diff --git a/tests/test_series.py b/tests/test_series.py index f76c2662d..61e8f9654 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1832,3 +1832,18 @@ def test_loc_callable() -> None: # GH 586 s = pd.Series([1, 2]) check(assert_type(s.loc[lambda x: x > 1], pd.Series), pd.Series) + + +def test_to_json_mode() -> None: + s = pd.Series([1, 2, 3, 4]) + result = s.to_json(orient="records", lines=True, mode="a") + result1 = s.to_json(orient="split", mode="w") + result2 = s.to_json(orient="table", mode="w") + check(assert_type(result, str), str) + check(assert_type(result1, str), str) + check(assert_type(result2, str), str) + if TYPE_CHECKING_INVALID_USAGE: + result3 = s.to_json(orient="records", lines=False, mode="a") # type: ignore[call-overload] + result4 = s.to_json(orient="records", mode="w") + check(assert_type(result3, Any), str) + check(assert_type(result4, str), str) From 5b44d703f87d51bfeff554414f44494cdea343c7 Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Wed, 10 May 2023 14:47:55 +0530 Subject: [PATCH 4/7] added warnings --- tests/test_frame.py | 4 +++- tests/test_series.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 4a8eccf2a..053c9e07a 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2593,7 +2593,9 @@ def test_to_json_mode() -> None: check(assert_type(result1, str), str) check(assert_type(result2, str), str) if TYPE_CHECKING_INVALID_USAGE: - result3 = df.to_json(orient="records", lines=False, mode="a") # type: ignore[call-overload] + result3 = df.to_json( + orient="records", lines=False, mode="a" + ) # pyright: ignore[reportGeneralTypeIssues] # type: ignore[call-overload] result4 = df.to_json(orient="records", mode="w") check(assert_type(result3, Any), str) check(assert_type(result4, str), str) diff --git a/tests/test_series.py b/tests/test_series.py index 61e8f9654..82d05931a 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1843,7 +1843,9 @@ def test_to_json_mode() -> None: check(assert_type(result1, str), str) check(assert_type(result2, str), str) if TYPE_CHECKING_INVALID_USAGE: - result3 = s.to_json(orient="records", lines=False, mode="a") # type: ignore[call-overload] + result3 = s.to_json( + orient="records", lines=False, mode="a" + ) # pyright: ignore[reportGeneralTypeIssues] # type: ignore[call-overload] result4 = s.to_json(orient="records", mode="w") check(assert_type(result3, Any), str) check(assert_type(result4, str), str) From 1c70c7a3250b43de36cef63c86418efbe5a79546 Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Wed, 10 May 2023 22:23:30 +0530 Subject: [PATCH 5/7] done required changes --- pandas-stubs/core/series.pyi | 28 ++++++++++++++-------------- tests/test_frame.py | 10 ++++------ tests/test_series.py | 10 ++++------ 3 files changed, 22 insertions(+), 26 deletions(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 62feaf9d1..d4446f8eb 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -488,72 +488,72 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def to_json( self, + *, path_or_buf: FilePath | WriteBuffer[str], - orient: JsonSeriesOrient | None = ..., + orient: Literal["records"], date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., force_ascii: _bool = ..., date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: _bool = ..., + lines: Literal[True], compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., - mode: Literal["w"] = ..., + mode: Literal["a"], ) -> None: ... @overload def to_json( self, + *, path_or_buf: None = ..., - orient: JsonSeriesOrient | None = ..., + orient: Literal["records"], date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., force_ascii: _bool = ..., date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: _bool = ..., + lines: Literal[True], compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., - mode: Literal["w"] = ..., + mode: Literal["a"], ) -> _str: ... @overload def to_json( self, - *, path_or_buf: FilePath | WriteBuffer[str], - orient: Literal["records"], + orient: JsonSeriesOrient | None = ..., date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., force_ascii: _bool = ..., date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: Literal[True], + lines: _bool = ..., compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., - mode: Literal["a"], + mode: Literal["w"] = ..., ) -> None: ... @overload def to_json( self, - *, path_or_buf: None = ..., - orient: Literal["records"], + orient: JsonSeriesOrient | None = ..., date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., force_ascii: _bool = ..., date_unit: Literal["s", "ms", "us", "ns"] = ..., default_handler: Callable[[Any], _str | float | _bool | list | dict] | None = ..., - lines: Literal[True], + lines: _bool = ..., compression: CompressionOptions = ..., index: _bool = ..., indent: int | None = ..., - mode: Literal["a"], + mode: Literal["w"] = ..., ) -> _str: ... def to_xarray(self) -> xr.DataArray: ... def items(self) -> Iterable[tuple[Hashable, S1]]: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 053c9e07a..18a2bee75 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2589,13 +2589,11 @@ def test_to_json_mode() -> None: result = df.to_json(orient="records", lines=True, mode="a") result1 = df.to_json(orient="split", mode="w") result2 = df.to_json(orient="columns", mode="w") + result4 = df.to_json(orient="records", mode="w") check(assert_type(result, str), str) check(assert_type(result1, str), str) check(assert_type(result2, str), str) + check(assert_type(result4, str), str) if TYPE_CHECKING_INVALID_USAGE: - result3 = df.to_json( - orient="records", lines=False, mode="a" - ) # pyright: ignore[reportGeneralTypeIssues] # type: ignore[call-overload] - result4 = df.to_json(orient="records", mode="w") - check(assert_type(result3, Any), str) - check(assert_type(result4, str), str) + result3 = df.to_json(orient="records", lines=False, mode="a") # type: ignore[call-overload] # pyright: ignore[reportGeneralTypeIssues] + assert_type(result3, Any) diff --git a/tests/test_series.py b/tests/test_series.py index 82d05931a..2b0bd0be0 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1839,13 +1839,11 @@ def test_to_json_mode() -> None: result = s.to_json(orient="records", lines=True, mode="a") result1 = s.to_json(orient="split", mode="w") result2 = s.to_json(orient="table", mode="w") + result4 = s.to_json(orient="records", mode="w") check(assert_type(result, str), str) check(assert_type(result1, str), str) check(assert_type(result2, str), str) + check(assert_type(result4, str), str) if TYPE_CHECKING_INVALID_USAGE: - result3 = s.to_json( - orient="records", lines=False, mode="a" - ) # pyright: ignore[reportGeneralTypeIssues] # type: ignore[call-overload] - result4 = s.to_json(orient="records", mode="w") - check(assert_type(result3, Any), str) - check(assert_type(result4, str), str) + result3 = s.to_json(orient="records", lines=False, mode="a") # type: ignore[call-overload] # pyright: ignore[reportGeneralTypeIssues] + assert_type(result3, Any) From eea6aa53e51eeb6e4a1e58e8b2ba8a53ca5c0fb7 Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Wed, 10 May 2023 22:24:08 +0530 Subject: [PATCH 6/7] Update series.pyi --- pandas-stubs/core/series.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index d4446f8eb..50216361c 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -488,8 +488,8 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def to_json( self, - *, path_or_buf: FilePath | WriteBuffer[str], + *, orient: Literal["records"], date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., @@ -506,8 +506,8 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def to_json( self, - *, path_or_buf: None = ..., + *, orient: Literal["records"], date_format: Literal["epoch", "iso"] | None = ..., double_precision: int = ..., From 42673179dfe044f637951d628fc70c81746bec9e Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Thu, 11 May 2023 10:16:39 +0530 Subject: [PATCH 7/7] removed the test --- tests/test_frame.py | 1 - tests/test_series.py | 1 - 2 files changed, 2 deletions(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 18a2bee75..6d6db78d9 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2596,4 +2596,3 @@ def test_to_json_mode() -> None: check(assert_type(result4, str), str) if TYPE_CHECKING_INVALID_USAGE: result3 = df.to_json(orient="records", lines=False, mode="a") # type: ignore[call-overload] # pyright: ignore[reportGeneralTypeIssues] - assert_type(result3, Any) diff --git a/tests/test_series.py b/tests/test_series.py index 2b0bd0be0..51118595d 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1846,4 +1846,3 @@ def test_to_json_mode() -> None: check(assert_type(result4, str), str) if TYPE_CHECKING_INVALID_USAGE: result3 = s.to_json(orient="records", lines=False, mode="a") # type: ignore[call-overload] # pyright: ignore[reportGeneralTypeIssues] - assert_type(result3, Any)