fix(python): ensure parametric testing cols=int definition respects allowed_dtypes

py-polars/polars/internals/construction.py

@@ -51,7 +51,6 @@
     _PANDAS_AVAILABLE,
     _PYARROW_AVAILABLE,
     _check_for_numpy,
-    _check_for_pandas,
 )
 from polars.dependencies import numpy as np
 from polars.dependencies import pandas as pd

py-polars/polars/testing/_parametric.py

@@ -559,8 +559,6 @@ def dataframes(
     │ 575050513 ┆ NaN        │
     └───────────┴────────────┘
     """  # noqa: 501
-    if isinstance(cols, int):
-        cols = columns(cols)
     if isinstance(min_size, int) and min_cols in (0, None):
         min_cols = 1
 
@@ -574,16 +572,16 @@ def dataframes(
     def draw_frames(draw: DrawFn) -> pli.DataFrame | pli.LazyFrame:
         with StringCache():
             # if not given, create 'n' cols with random dtypes
-            if cols is None:
-                n = between(
+            if cols is None or isinstance(cols, int):
+                n = cols or between(
                     draw, int, min_=(min_cols or 0), max_=(max_cols or MAX_COLS)
                 )
                 dtypes_ = [draw(sampled_from(selectable_dtypes)) for _ in range(n)]
                 coldefs = columns(cols=n, dtype=dtypes_)
             elif isinstance(cols, column):
                 coldefs = [cols]
             else:
-                coldefs = list(cols)  # type: ignore[arg-type]
+                coldefs = list(cols)
 
             # append any explicitly provided cols
             coldefs.extend(include_cols or ())

py-polars/tests/parametric/test_dataframe.py

@@ -18,6 +18,13 @@ def test_repr(df: pl.DataFrame) -> None:
     assert_frame_equal(df, df, check_exact=True, nans_compare_equal=True)
 
 
+@given(df=dataframes(cols=10, max_size=1, allowed_dtypes=[pl.UInt8, pl.Int8]))
+@settings(max_examples=3)
+def test_dtype_integer_cols(df: pl.DataFrame) -> None:
+    # ensure dtype constraint works in conjunction with 'n' cols
+    assert all(tp in (pl.UInt8, pl.Int8) for tp in df.schema.values())
+
+
 @given(
     df=dataframes(
         min_size=1, min_cols=1, null_probability=0.25, excluded_dtypes=[pl.Utf8]