Add type annotation for projection variable (#2100)

py-polars/polars/internals/frame.py

@@ -58,7 +58,7 @@
 
 from polars._html import NotebookFormatter
 from polars.datatypes import Boolean, DataType, Datetime, UInt32, py_type_to_dtype
-from polars.utils import _process_null_values
+from polars.utils import _process_null_values, is_int_sequence, is_str_sequence
 
 try:
     import pandas as pd
@@ -463,16 +463,15 @@ def read_csv(
             if isinstance(file, StringIO):
                 file = file.getvalue().encode()
 
-        projection: Optional[List[int]] = None
+        projection: Optional[Sequence[int]] = None
         if columns:
-            if isinstance(columns, list):
-                if all(isinstance(i, int) for i in columns):
-                    projection = columns  # type: ignore
-                    columns = None
-                elif not all(isinstance(i, str) for i in columns):
-                    raise ValueError(
-                        "columns arg should contain a list of all integers or all strings values."
-                    )
+            if is_int_sequence(columns):
+                projection = columns
+                columns = None
+            elif not is_str_sequence(columns):
+                raise ValueError(
+                    "columns arg should contain a list of all integers or all strings values."
+                )
 
         dtype_list: Optional[List[Tuple[str, Type[DataType]]]] = None
         dtype_slice: Optional[List[Type[DataType]]] = None
@@ -531,16 +530,15 @@ def read_parquet(
         n_rows
             Stop reading from parquet file after reading ``n_rows``.
         """
-        projection: Optional[List[int]] = None
+        projection: Optional[Sequence[int]] = None
         if columns:
-            if isinstance(columns, list):
-                if all(isinstance(i, int) for i in columns):
-                    projection = columns  # type: ignore
-                    columns = None
-                elif not all(isinstance(i, str) for i in columns):
-                    raise ValueError(
-                        "columns arg should contain a list of all integers or all strings values."
-                    )
+            if is_int_sequence(columns):
+                projection = columns
+                columns = None
+            elif not is_str_sequence(columns):
+                raise ValueError(
+                    "columns arg should contain a list of all integers or all strings values."
+                )
 
         self = DataFrame.__new__(DataFrame)
         self._df = PyDataFrame.read_parquet(file, columns, projection, n_rows)
@@ -568,16 +566,15 @@ def read_ipc(
         -------
         DataFrame
         """
-        projection: Optional[List[int]] = None
+        projection: Optional[Sequence[int]] = None
         if columns:
-            if isinstance(columns, list):
-                if all(isinstance(i, int) for i in columns):
-                    projection = columns  # type: ignore
-                    columns = None
-                elif not all(isinstance(i, str) for i in columns):
-                    raise ValueError(
-                        "columns arg should contain a list of all integers or all strings values."
-                    )
+            if is_int_sequence(columns):
+                projection = columns
+                columns = None
+            elif not is_str_sequence(columns):
+                raise ValueError(
+                    "columns arg should contain a list of all integers or all strings values."
+                )
 
         self = DataFrame.__new__(DataFrame)
         self._df = PyDataFrame.read_ipc(file, columns, projection, n_rows)

py-polars/polars/utils.py

@@ -1,9 +1,15 @@
 import ctypes
+import sys
 from datetime import date, datetime, timedelta, timezone
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Dict, Iterable, List, Sequence, Tuple, Type, Union
 
 import numpy as np
 
+if sys.version_info >= (3, 10):
+    from typing import TypeGuard
+else:
+    from typing_extensions import TypeGuard
+
 
 def _process_null_values(
     null_values: Union[None, str, List[str], Dict[str, str]] = None,
@@ -51,3 +57,23 @@ def _datetime_to_pl_timestamp(dt: datetime) -> int:
 def _date_to_pl_date(d: date) -> int:
     dt = datetime.combine(d, datetime.min.time()).replace(tzinfo=timezone.utc)
     return int(dt.timestamp()) // (3600 * 24)
+
+
+def is_str_sequence(
+    val: Sequence[object], allow_str: bool = False
+) -> TypeGuard[Sequence[str]]:
+    """
+    Checks that `val` is a sequence of strings. Note that a single string is a sequence of strings
+    by definition, use `allow_str=False` to return False on a single string
+    """
+    if (not allow_str) and isinstance(val, str):
+        return False
+    return _is_iterable_of(val, Sequence, str)
+
+
+def is_int_sequence(val: Sequence[object]) -> TypeGuard[Sequence[int]]:
+    return _is_iterable_of(val, Sequence, int)
+
+
+def _is_iterable_of(val: Iterable, itertype: Type, eltype: Type) -> bool:
+    return isinstance(val, itertype) and all(isinstance(x, eltype) for x in val)