From 9cb7e8913531d26e9157ad3642c7c091101d3bd2 Mon Sep 17 00:00:00 2001 From: Erik-Jan van Kesteren Date: Fri, 23 Feb 2024 14:14:56 +0100 Subject: [PATCH] Make mypy happy --- metasyncontrib/disclosure/discrete.py | 4 ++-- metasyncontrib/disclosure/numerical.py | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/metasyncontrib/disclosure/discrete.py b/metasyncontrib/disclosure/discrete.py index cb03ffe..10d6352 100644 --- a/metasyncontrib/disclosure/discrete.py +++ b/metasyncontrib/disclosure/discrete.py @@ -25,8 +25,8 @@ class DisclosureUniqueKey(UniqueKeyDistribution): @classmethod def _fit(cls, values: pl.Series, n_avg: int = 11): orig_dist = super()._fit(values) - if orig_dist.consecutive == 1: - return cls(0, 1) + if orig_dist.consecutive: + return cls(0, True) sub_values = micro_aggregate(values, n_avg) return super()._fit(sub_values) diff --git a/metasyncontrib/disclosure/numerical.py b/metasyncontrib/disclosure/numerical.py index 142ac82..a79b581 100644 --- a/metasyncontrib/disclosure/numerical.py +++ b/metasyncontrib/disclosure/numerical.py @@ -14,8 +14,7 @@ class DisclosureNumerical(BaseDistribution): """Class for numerical distributions of the disclosure kind.""" @classmethod - def fit(cls, series: Union[Sequence, pl.Series], *args, - n_avg: int = 11, **kwargs) -> BaseDistribution: + def fit(cls, series, *args, n_avg: int = 11, **kwargs) -> BaseDistribution: pl_series = cls._to_series(series) sub_series = micro_aggregate(pl_series, n_avg) return cls._fit(sub_series, *args, **kwargs)