improve type inference

Author: JohnMount

data_algebra/util.py

@@ -75,16 +75,21 @@ def guess_carried_scalar_type(col):
     :param col: column or scalar to inspect
     :return: type of first non-None entry, if any , else type(None)
     """
+    # check for scalars first
     ct = map_type_to_canonical(type(col))
     if ct in {str, int, float, bool, type(None), numpy.int64, numpy.float64,
               datetime.datetime, datetime.date, datetime.timedelta}:
         return ct
+    # look at a list or Series
+    if isinstance(col, data_algebra.default_data_model.pd.core.series.Series):
+        col = col.values
     if len(col) < 1:
         return type(None)
-    idx = col.notna().idxmax()
-    if idx is None:
-        return map_type_to_canonical(type(col[0]))
-    return map_type_to_canonical(type(col[idx]))
+    good_idx = numpy.where(numpy.logical_not(data_algebra.default_data_model.pd.isna(col)))[0]
+    test_idx = 0
+    if len(good_idx) > 0:
+        test_idx = good_idx[0]
+    return map_type_to_canonical(type(col[test_idx]))
 
 
 def guess_column_types(d, *, columns=None):
@@ -106,7 +111,7 @@ def guess_column_types(d, *, columns=None):
     res = dict()
     for c in columns:
         gt = guess_carried_scalar_type(d[c])
-        if (gt is None) or (not isinstance(gt, type)) or str(gt).endswith('.Series\'>'):
+        if (gt is None) or (not isinstance(gt, type)) or gt == data_algebra.default_data_model.pd.core.series.Series:
             # pandas.concat() poisons types with Series, don't allow that
             return dict()
         res[c] = gt

tests/test_type_check_problem.py

@@ -1,13 +1,9 @@
 
-import numpy
-
 import data_algebra
 from data_algebra.data_ops import *
 import data_algebra.util
 import data_algebra.test_util
 
-import pytest
-
 
 def test_type_check_problem_1():
     d = data_algebra.default_data_model.pd.DataFrame({
@@ -16,8 +12,13 @@ def test_type_check_problem_1():
     assert not isinstance(d['x'][0], data_algebra.default_data_model.pd.core.series.Series)
     d2 = data_algebra.default_data_model.pd.concat([d, d])
     # pandas concat mucks columns into series (reset_index() fixes that)
+    # check that we see that problem
     assert isinstance(d2['x'][0], data_algebra.default_data_model.pd.core.series.Series)
+    # reset index can clear the problem
     d3 = d2.reset_index(drop=True, inplace=False)
     assert not isinstance(d3['x'][0], data_algebra.default_data_model.pd.core.series.Series)
+    # see what type inspection we get
     td2 = describe_table(d2, table_name='d2')
-    assert (td2.column_types is None) or (len(td2.column_types) == 0)  # don't pick up Series types!
+    td3 = describe_table(d2, table_name='d3')
+    assert td3.column_types['x'] == int
+    assert td2.column_types == td3.column_types  # don't pick up Series types!