address more comments

lithomas1 · Jul 1, 2024 · 60287e1 · 60287e1
1 parent 25c25d4
commit 60287e1
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/python/cudf/cudf/_lib/pylibcudf/io/types.pyx b/python/cudf/cudf/_lib/pylibcudf/io/types.pyx
@@ -205,15 +205,15 @@ cdef class SinkInfo:
         if not sinks:
             raise ValueError("Need to pass at least one sink")
 
+        if isinstance(sinks[0], os.PathLike):
+            sinks = [os.path.expanduser(s) for s in sinks]
+
         cdef object initial_sink_cls = type(sinks[0])
 
         if not all(isinstance(s, initial_sink_cls) for s in sinks):
             raise ValueError("All sinks must be of the same type!")
 
-        if isinstance(sinks[0], os.PathLike):
-            sinks = [os.path.expanduser(s) for s in sinks]
-
-        if isinstance(sinks[0], (io.StringIO, io.BytesIO, io.TextIOBase)):
+        if initial_sink_cls in {io.StringIO, io.BytesIO, io.TextIOBase}:
             data_sinks.reserve(len(sinks))
             if isinstance(sinks[0], (io.StringIO, io.BytesIO)):
                 for s in sinks:

diff --git a/python/cudf/cudf/pylibcudf_tests/io/test_json.py b/python/cudf/cudf/pylibcudf_tests/io/test_json.py
@@ -14,12 +14,11 @@ def test_write_json_basic(table_data, source_or_sink, lines, rows_per_chunk):
     plc_table_w_meta, pa_table = table_data
     sink = source_or_sink
 
-    kwargs = {}
-    if rows_per_chunk <= plc_table_w_meta.tbl.num_rows():
-        kwargs["rows_per_chunk"] = rows_per_chunk
-
     plc.io.json.write_json(
-        plc.io.SinkInfo([sink]), plc_table_w_meta, lines=lines, **kwargs
+        plc.io.SinkInfo([sink]),
+        plc_table_w_meta,
+        lines=lines,
+        rows_per_chunk=rows_per_chunk,
     )
 
     exp = pa_table.to_pandas()