|
2 | 2 |
|
3 | 3 | import inspect |
4 | 4 | import logging |
5 | | -from typing import Any, Callable, Optional, Union |
6 | | - |
7 | | -import pandas as pd |
8 | | -import polars as pl |
9 | | - |
10 | | -# Import fully qualified types to satisfy disallow_any_unimported |
11 | | -from pandas import DataFrame as PandasDataFrame |
12 | | -from polars import DataFrame as PolarsDataFrame |
| 5 | +from typing import TYPE_CHECKING, Any, Callable, Optional, Union |
| 6 | + |
| 7 | +# Lazy imports - only import what's available |
| 8 | +try: |
| 9 | + import pandas as pd |
| 10 | + from pandas import DataFrame as PandasDataFrame |
| 11 | + |
| 12 | + HAS_PANDAS = True |
| 13 | +except ImportError: # pragma: no cover |
| 14 | + pd = None # type: ignore |
| 15 | + PandasDataFrame = None # type: ignore |
| 16 | + HAS_PANDAS = False |
| 17 | + |
| 18 | +try: |
| 19 | + import polars as pl |
| 20 | + from polars import DataFrame as PolarsDataFrame |
| 21 | + |
| 22 | + HAS_POLARS = True |
| 23 | +except ImportError: # pragma: no cover |
| 24 | + pl = None # type: ignore |
| 25 | + PolarsDataFrame = None # type: ignore |
| 26 | + HAS_POLARS = False |
| 27 | + |
| 28 | +# Build DataFrame type dynamically based on what's available |
| 29 | +if TYPE_CHECKING: |
| 30 | + # For static type checking, assume both are available |
| 31 | + from pandas import DataFrame as PandasDataFrame |
| 32 | + from polars import DataFrame as PolarsDataFrame |
| 33 | + |
| 34 | + DataFrameType = Union[PandasDataFrame, PolarsDataFrame] |
| 35 | +else: |
| 36 | + # For runtime, build type tuple from available libraries |
| 37 | + _available_types = [] |
| 38 | + if HAS_PANDAS: |
| 39 | + _available_types.append(PandasDataFrame) |
| 40 | + if HAS_POLARS: |
| 41 | + _available_types.append(PolarsDataFrame) |
| 42 | + |
| 43 | + if not _available_types: # pragma: no cover |
| 44 | + raise ImportError( |
| 45 | + "No DataFrame library found. Please install Pandas or Polars: pip install pandas OR pip install polars" |
| 46 | + ) |
13 | 47 |
|
14 | | -DataFrameType = Union[PandasDataFrame, PolarsDataFrame] |
| 48 | + DataFrameType = Union[tuple(_available_types)] |
15 | 49 |
|
16 | 50 |
|
17 | 51 | def assert_is_dataframe(obj: Any, context: str) -> None: |
18 | | - if not isinstance(obj, (pd.DataFrame, pl.DataFrame)): |
19 | | - raise AssertionError(f"Wrong {context}. Expected DataFrame, got {type(obj).__name__} instead.") |
| 52 | + # Build type tuple dynamically based on available libraries |
| 53 | + dataframe_types: list[Any] = [] |
| 54 | + if HAS_PANDAS and pd is not None: |
| 55 | + dataframe_types.append(pd.DataFrame) |
| 56 | + if HAS_POLARS and pl is not None: |
| 57 | + dataframe_types.append(pl.DataFrame) |
| 58 | + |
| 59 | + if not isinstance(obj, tuple(dataframe_types)): |
| 60 | + available_libs = [] |
| 61 | + if HAS_PANDAS: |
| 62 | + available_libs.append("Pandas") |
| 63 | + if HAS_POLARS: |
| 64 | + available_libs.append("Polars") |
| 65 | + libs_str = " or ".join(available_libs) |
| 66 | + raise AssertionError(f"Wrong {context}. Expected {libs_str} DataFrame, got {type(obj).__name__} instead.") |
20 | 67 |
|
21 | 68 |
|
22 | 69 | def format_param_context( |
@@ -64,21 +111,35 @@ def get_parameter_name( |
64 | 111 | def describe_dataframe(df: DataFrameType, include_dtypes: bool = False) -> str: |
65 | 112 | result = f"columns: {list(df.columns)}" |
66 | 113 | if include_dtypes: |
67 | | - if isinstance(df, pd.DataFrame): |
| 114 | + if HAS_PANDAS and pd is not None and isinstance(df, pd.DataFrame): |
68 | 115 | readable_dtypes = [dtype.name for dtype in df.dtypes] |
69 | 116 | result += f" with dtypes {readable_dtypes}" |
70 | | - else: |
| 117 | + elif HAS_POLARS and pl is not None and isinstance(df, pl.DataFrame): |
71 | 118 | result += f" with dtypes {df.dtypes}" |
72 | 119 | return result |
73 | 120 |
|
74 | 121 |
|
75 | 122 | def log_dataframe_input(level: int, func_name: str, df: Any, include_dtypes: bool) -> None: |
76 | | - if isinstance(df, (pd.DataFrame, pl.DataFrame)): |
| 123 | + # Build type tuple dynamically based on available libraries |
| 124 | + dataframe_types: list[Any] = [] |
| 125 | + if HAS_PANDAS and pd is not None: |
| 126 | + dataframe_types.append(pd.DataFrame) |
| 127 | + if HAS_POLARS and pl is not None: |
| 128 | + dataframe_types.append(pl.DataFrame) |
| 129 | + |
| 130 | + if isinstance(df, tuple(dataframe_types)): |
77 | 131 | logging.log( |
78 | 132 | level, f"Function {func_name} parameters contained a DataFrame: {describe_dataframe(df, include_dtypes)}" |
79 | 133 | ) |
80 | 134 |
|
81 | 135 |
|
82 | 136 | def log_dataframe_output(level: int, func_name: str, df: Any, include_dtypes: bool) -> None: |
83 | | - if isinstance(df, (pd.DataFrame, pl.DataFrame)): |
| 137 | + # Build type tuple dynamically based on available libraries |
| 138 | + dataframe_types: list[Any] = [] |
| 139 | + if HAS_PANDAS and pd is not None: |
| 140 | + dataframe_types.append(pd.DataFrame) |
| 141 | + if HAS_POLARS and pl is not None: |
| 142 | + dataframe_types.append(pl.DataFrame) |
| 143 | + |
| 144 | + if isinstance(df, tuple(dataframe_types)): |
84 | 145 | logging.log(level, f"Function {func_name} returned a DataFrame: {describe_dataframe(df, include_dtypes)}") |
0 commit comments