Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
chore: clean up OfType selector
Browse files Browse the repository at this point in the history
cpcloud committed Aug 25, 2024

Verified

This commit was signed with the committer’s verified signature.
oddgrd Oddbjørn Grødem
1 parent a116af3 commit ce9d31d
Showing 1 changed file with 18 additions and 6 deletions.
24 changes: 18 additions & 6 deletions ibis/selectors.py
Original file line number Diff line number Diff line change
@@ -128,6 +128,15 @@ def numeric() -> Selector:
return of_type(dt.Numeric)


class OfType(Selector):
predicate: Callable[[dt.DataType], bool]

def expand_names(self, table: ir.Table) -> frozenset[str]:
return frozenset(
name for name, typ in table.schema().items() if self.predicate(typ)
)


@public
def of_type(dtype: dt.DataType | str | type[dt.DataType]) -> Selector:
"""Select columns of type `dtype`.
@@ -186,17 +195,20 @@ def of_type(dtype: dt.DataType | str | type[dt.DataType]) -> Selector:
"struct": dt.Struct,
"temporal": dt.Temporal,
}
if cls := abstract.get(dtype.lower()):
predicate = lambda col: isinstance(col.type(), cls)

if dtype_cls := abstract.get(dtype.lower()):
predicate = lambda typ, dtype_cls=dtype_cls: isinstance(typ, dtype_cls)
else:
dtype = dt.dtype(dtype)
predicate = lambda col: col.type() == dtype
predicate = lambda typ, dtype=dtype: typ == dtype

elif inspect.isclass(dtype) and issubclass(dtype, dt.DataType):
predicate = lambda col: isinstance(col.type(), dtype)
predicate = lambda typ, dtype_cls=dtype: isinstance(typ, dtype_cls)
else:
dtype = dt.dtype(dtype)
predicate = lambda col: col.type() == dtype
return where(predicate)
predicate = lambda typ, dtype=dtype: typ == dtype

return OfType(predicate)


class StartsWith(Selector):

0 comments on commit ce9d31d

Please sign in to comment.