fix tests

sdv-dev · Oct 21, 2019 · a718fd5 · a718fd5
1 parent 0010035
commit a718fd5
Show file tree

Hide file tree

Showing 5 changed files with 199 additions and 155 deletions.
diff --git a/sdv/modeler.py b/sdv/modeler.py
@@ -255,6 +255,8 @@ def _get_extensions(self, parent, children, tables):
                 Names of the children.
             tables (dict):
                 Previously processed tables.
+        Returns:
+            pandas.DataFrame
         """
         extensions = []
 
@@ -293,7 +295,7 @@ def cpa(self, table_name, tables):
                 previously processed tables.
 
         Returns:
-            None
+            pandas.DataFrame
         """
         LOGGER.info('Modeling %s', table_name)
 

diff --git a/sdv/sdv.py b/sdv/sdv.py
@@ -74,8 +74,11 @@ def sample_rows(self, table_name, num_rows, sample_children=True, reset_primary_
             raise NotFittedError('SDV instance has not been fitted')
 
         return self.sampler.sample_rows(
-            table_name, num_rows, sample_children=sample_children,
-            reset_primary_keys=reset_primary_keys)
+            table_name,
+            num_rows,
+            sample_children=sample_children,
+            reset_primary_keys=reset_primary_keys
+        )
 
     def sample_table(self, table_name, num_rows=None, reset_primary_keys=False):
         """Samples the given table to its original size.