Merge pull request #334 from taylorfturner/sam_comments_fix

Quick Fixes
capitalone · Aug 28, 2023 · 5be1ed8 · 5be1ed8
2 parents b9e5c56 + fab0fc3
commit 5be1ed8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/synthetic_data/distinct_generators/datetime_generator.py b/synthetic_data/distinct_generators/datetime_generator.py
@@ -1,6 +1,6 @@
 """Contains a datetime generator."""
 from datetime import datetime
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 import pandas as pd
@@ -43,7 +43,7 @@ def generate_datetime(
 
 def random_datetimes(
     rng: Generator,
-    format: Optional[str] = None,
+    format: Optional[List[str]] = None,
     min: pd.Timestamp = None,
     max: pd.Timestamp = None,
     num_rows: int = 1,
@@ -70,6 +70,9 @@ def random_datetimes(
     if not format:
         format = ["%B %d %Y %H:%M:%S"]
 
+    if not isinstance(format, list):
+        raise Exception("format must be of type `list`")
+
     for i in range(num_rows):
         date_format = rng.choice(format)
         datetime = generate_datetime(

diff --git a/synthetic_data/generators.py b/synthetic_data/generators.py
@@ -115,7 +115,7 @@ def _generate_uncorrelated_column_data(self, num_samples):
             col_["rng"] = self.rng
             col_["num_rows"] = num_samples
 
-            if (generator_name == "string") or (generator_name == "text"):
+            if generator_name in ["string", "text"]:
                 if col_.get("categorical", False):
                     generator_name = "categorical"
                     total = 0