Fix issue with polars update (#30)

examples/tutorial.ipynb

@@ -22,7 +22,6 @@
     "\n",
     "import numpy as np\n",
     "import polars as pl\n",
-    "from matplotlib import pyplot as plt\n",
     "from metasyn import MetaFrame, demo_file\n",
     "from metasyn.provider import DistributionProviderList\n",
     "from metasyn.distribution import (\n",
@@ -320,6 +319,17 @@
     "Below we will look at what happens to our parameters if we add a single new value (outlier) to the data. We do this for both the base metasyn implementation and the disclosure control implementation. We expect that following the rules of disclosure control, a single outlier should have a smaller (and limited) effect on the results than with the base metasyn implementation."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "106e5463",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# If you get an error, install matplotlib with `pip install matplotlib`\n",
+    "from matplotlib import pyplot as plt\n"
+   ]
+  },
   {
    "attachments": {},
    "cell_type": "markdown",
@@ -682,7 +692,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.3"
+   "version": "3.11.6"
   },
   "vscode": {
    "interpreter": {

metasyncontrib/disclosure/utils.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import datetime as dt
 from typing import NamedTuple, Optional
 
 import numpy as np
@@ -61,11 +60,7 @@ def _create_subsample( # pylint: disable=too-many-locals
         sub_values = []
         for block in block_values:
             mean_time = pl.Series(block).dt.cast_time_unit("us").mean()
-            assert mean_time is not None
-            sec_since_1970 = mean_time / 1e6
-            sub_values.append(
-                dt.datetime.utcfromtimestamp(0) + dt.timedelta(seconds=sec_since_1970)
-            )
+            sub_values.append(mean_time)
     return sub_values, dominance
 
 

tests/test_other_dist.py

@@ -1,3 +1,5 @@
+import datetime as dt
+
 import numpy as np
 import polars as pl
 from metasyn.distribution.categorical import MultinoulliDistribution
@@ -14,20 +16,21 @@
 
 
 @mark.parametrize(
-    "class_norm,class_disc",
+    "class_norm,class_disc,input_type",
     [
-        (DateUniformDistribution, DisclosureDate),
-        (DateTimeUniformDistribution, DisclosureDateTime),
-        (TimeUniformDistribution, DisclosureTime),
+        (DateUniformDistribution, DisclosureDate, dt.date),
+        (DateTimeUniformDistribution, DisclosureDateTime, dt.datetime),
+        (TimeUniformDistribution, DisclosureTime, dt.time),
     ],
 )
-def test_datetime(class_norm, class_disc):
+def test_datetime(class_norm, class_disc, input_type):
     dist_norm = class_norm.default_distribution()
     series = pl.Series([dist_norm.draw() for _ in range(100)])
     dist_norm = class_norm.fit(series)
     dist_disc = class_disc.fit(series)
     assert dist_norm.lower < dist_disc.lower
     assert dist_norm.upper > dist_disc.upper
+    assert isinstance(dist_norm.draw(), input_type)
     if not isinstance(dist_norm, DateUniformDistribution):
         assert dist_norm.precision == dist_disc.precision