Skip to content

Commit

Permalink
Replaced callable check
Browse files Browse the repository at this point in the history
Signed-off-by: Elena Khaustova <[email protected]>
  • Loading branch information
ElenaKhaustova committed Jan 7, 2025
1 parent 87d5e62 commit aaf4a72
Showing 1 changed file with 6 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,11 @@ def _grandparent(path: str) -> str:
return str(grandparent)


def _islambda(obj: object):
"""Check if object is a lambda function."""
return callable(obj) and hasattr(obj, "__name__") and obj.__name__ == "<lambda>"


class PartitionedDataset(AbstractDataset[dict[str, Any], dict[str, Callable[[], Any]]]):
"""``PartitionedDataset`` loads and saves partitioned file-like data using the
underlying dataset definition. For filesystem level operations it uses `fsspec`:
Expand Down Expand Up @@ -311,7 +316,7 @@ def save(self, data: dict[str, Any]) -> None:
# join the protocol back since tools like PySpark may rely on it
kwargs[self._filepath_arg] = self._join_protocol(partition)
dataset = self._dataset_type(**kwargs) # type: ignore
if callable(partition_data):
if _islambda(partition_data):
partition_data = partition_data() # noqa: PLW2901
dataset.save(partition_data)
self._invalidate_caches()
Expand Down

0 comments on commit aaf4a72

Please sign in to comment.