Specify maxsplit to handle colon symbols properly (#247)

google · Aug 15, 2023 · 613ba71 · 613ba71
1 parent a069d89
commit 613ba71
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/scripts/build_model.py b/scripts/build_model.py
@@ -38,7 +38,7 @@ def aggregate_scores(
     if not row:
       continue
     feature = row.split('\t')[0]
-    feature_group, feature_content = feature.split(':')
+    feature_group, feature_content = feature.split(':', 1)
     score = float(row.split('\t')[1])
     decision_trees.setdefault(feature_group, {})
     decision_trees[feature_group].setdefault(feature_content, 0)

diff --git a/scripts/tests/test_build_model.py b/scripts/tests/test_build_model.py
@@ -57,6 +57,14 @@ def test_blank_line(self) -> None:
         }
     }, 'should skip blank lines.')
 
+  def test_colon(self) -> None:
+    weights = ['AB::\t8.123']
+    model = build_model.aggregate_scores(weights)
+    self.assertDictEqual(
+        model, {'AB': {
+            ':': 8.123
+        }}, 'should consider the first colon only as a delimiter.')
+
 
 class TestRoundModel(unittest.TestCase):