Skip to content

Commit

Permalink
Added safeguard to avoid confusion when calling as_dataclas/attrs
Browse files Browse the repository at this point in the history
  • Loading branch information
AKuederle committed Dec 15, 2023
1 parent e6a7032 commit 034e024
Showing 1 changed file with 9 additions and 4 deletions.
13 changes: 9 additions & 4 deletions tpcp/_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -655,9 +655,12 @@ def __init__(
self.groupby_cols = groupby_cols
self.subset_index = subset_index

@staticmethod
def as_dataclass():
@classmethod
def as_dataclass(cls):
"""Return a version of the Dataset class that can be subclassed using dataclasses."""
# You are only allowed to call this method on the base class not the subclass!
assert cls is Dataset, "You can only call `as_dataclass` on the base class `Dataset`!"

import dataclasses # pylint: disable=import-outside-toplevel

@dataclasses.dataclass(eq=False, repr=False, order=False)
Expand All @@ -669,12 +672,14 @@ class DatasetDc(_Dataset):

return DatasetDc

@staticmethod
def as_attrs():
@classmethod
def as_attrs(cls):
"""Return a version of the Dataset class that can be subclassed using `attrs` defined classes.
Note, this requires `attrs` to be installed!
"""
assert cls is Dataset, "You can only call `as_attrs` on the base class `Dataset`!"

from attrs import define # pylint: disable=import-outside-toplevel

@define(eq=False, repr=False, order=False, kw_only=True, slots=False)
Expand Down

0 comments on commit 034e024

Please sign in to comment.