|
55 | 55 | ColspaceArgType,
|
56 | 56 | CompressionOptions,
|
57 | 57 | Dtype,
|
| 58 | + DtypeObj, |
58 | 59 | FilePathOrBuffer,
|
59 | 60 | FillnaOptions,
|
60 | 61 | FloatFormatType,
|
@@ -4297,50 +4298,25 @@ def check_int_infer_dtype(dtypes):
|
4297 | 4298 | if not include.isdisjoint(exclude):
|
4298 | 4299 | raise ValueError(f"include and exclude overlap on {(include & exclude)}")
|
4299 | 4300 |
|
4300 |
| - # We raise when both include and exclude are empty |
4301 |
| - # Hence, we can just shrink the columns we want to keep |
4302 |
| - keep_these = np.full(self.shape[1], True) |
4303 |
| - |
4304 |
| - def extract_unique_dtypes_from_dtypes_set( |
4305 |
| - dtypes_set: frozenset[Dtype], unique_dtypes: np.ndarray |
4306 |
| - ) -> list[Dtype]: |
4307 |
| - extracted_dtypes = [ |
4308 |
| - unique_dtype |
4309 |
| - for unique_dtype in unique_dtypes |
4310 |
| - if ( |
4311 |
| - issubclass( |
4312 |
| - # error: Argument 1 to "tuple" has incompatible type |
4313 |
| - # "FrozenSet[Union[ExtensionDtype, Union[str, Any], Type[str], |
4314 |
| - # Type[float], Type[int], Type[complex], Type[bool], |
4315 |
| - # Type[object]]]"; expected "Iterable[Union[type, Tuple[Any, |
4316 |
| - # ...]]]" |
4317 |
| - unique_dtype.type, |
4318 |
| - tuple(dtypes_set), # type: ignore[arg-type] |
4319 |
| - ) |
4320 |
| - or ( |
4321 |
| - np.number in dtypes_set |
4322 |
| - and getattr(unique_dtype, "_is_numeric", False) |
4323 |
| - ) |
4324 |
| - ) |
4325 |
| - ] |
4326 |
| - return extracted_dtypes |
| 4301 | + def dtype_predicate(dtype: DtypeObj, dtypes_set) -> bool: |
| 4302 | + return issubclass(dtype.type, tuple(dtypes_set)) or ( |
| 4303 | + np.number in dtypes_set and getattr(dtype, "_is_numeric", False) |
| 4304 | + ) |
4327 | 4305 |
|
4328 |
| - unique_dtypes = self.dtypes.unique() |
| 4306 | + def predicate(arr: ArrayLike) -> bool: |
| 4307 | + dtype = arr.dtype |
| 4308 | + if include: |
| 4309 | + if not dtype_predicate(dtype, include): |
| 4310 | + return False |
4329 | 4311 |
|
4330 |
| - if include: |
4331 |
| - included_dtypes = extract_unique_dtypes_from_dtypes_set( |
4332 |
| - include, unique_dtypes |
4333 |
| - ) |
4334 |
| - keep_these &= self.dtypes.isin(included_dtypes) |
| 4312 | + if exclude: |
| 4313 | + if dtype_predicate(dtype, exclude): |
| 4314 | + return False |
4335 | 4315 |
|
4336 |
| - if exclude: |
4337 |
| - excluded_dtypes = extract_unique_dtypes_from_dtypes_set( |
4338 |
| - exclude, unique_dtypes |
4339 |
| - ) |
4340 |
| - keep_these &= ~self.dtypes.isin(excluded_dtypes) |
| 4316 | + return True |
4341 | 4317 |
|
4342 |
| - # error: "ndarray" has no attribute "values" |
4343 |
| - return self.iloc[:, keep_these.values] # type: ignore[attr-defined] |
| 4318 | + mgr = self._mgr._get_data_subset(predicate) |
| 4319 | + return type(self)(mgr).__finalize__(self) |
4344 | 4320 |
|
4345 | 4321 | def insert(self, loc, column, value, allow_duplicates: bool = False) -> None:
|
4346 | 4322 | """
|
|
0 commit comments