fix bug where infinity could be extracted as number

ecohealthalliance · Apr 26, 2019 · 8c1ad2c · 8c1ad2c
1 parent 2623431
commit 8c1ad2c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/epitator/raw_number_annotator.py b/epitator/raw_number_annotator.py
@@ -6,6 +6,7 @@
 from .date_annotator import DateAnnotator
 from . import utils
 import re
+import math
 
 
 def is_valid_number(num_string):
@@ -15,7 +16,7 @@ def is_valid_number(num_string):
     if num_string[0] == '0' and len(num_string) > 1:
         return False
     value = utils.parse_spelled_number(num_string)
-    return value is not None
+    return value not in [None, math.inf]
 
 
 class RawNumberAnnotator(Annotator):

diff --git a/tests/annotator/test_count_annotator.py b/tests/annotator/test_count_annotator.py
@@ -361,6 +361,10 @@ def test_count_list(self):
                          if 'case' in count.metadata['attributes']]
         self.assertSequenceEqual(actual_counts, expected_counts)
 
+    def test_infinity(self):
+        doc = AnnoDoc("Nmult 22-InF cases of Ebola.")
+        doc.add_tier(self.annotator)
+
 
 if __name__ == '__main__':
     unittest.main()