Make mypy happy

sodascience · Feb 23, 2024 · 9cb7e89 · 9cb7e89
1 parent 9bbdac6
commit 9cb7e89
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git 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
@@ -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)