Pass allow_bare callable to CityKeyIter.

mit-ccc · Aug 26, 2018 · 4465a93 · 4465a93
1 parent c26b9a9
commit 4465a93
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/litecoder/usa.py b/litecoder/usa.py
@@ -89,8 +89,8 @@ def __call__(self, row, name):
 
 class CityKeyIter:
 
-    def __init__(self, *args, **kwargs):
-        self.allow_bare = AllowBareCityName(*args, **kwargs)
+    def __init__(self, allow_bare):
+        self.allow_bare = allow_bare
 
     def __call__(self, row):
         """Enumerate index keys for a city.
@@ -245,7 +245,9 @@ def __init__(self, bare_name_blocklist=None):
     def build(self):
         """Index all US cities.
         """
-        iter_keys = CityKeyIter(blocklist=self.bare_name_blocklist)
+        allow_bare = AllowBareCityName(blocklist=self.bare_name_blocklist)
+
+        iter_keys = CityKeyIter(allow_bare)
 
         # Deduped cities.
         cities = WOFLocality.clean_us_cities()

diff --git a/tests/prod_db/test_bare_name_blocklist.py b/tests/prod_db/test_bare_name_blocklist.py
@@ -5,6 +5,7 @@
 from litecoder.usa import USCityIndex
 
 
+# TODO: Possible to speed this up?
 def test_bare_name_bloclist():
     """Blocklisted bare names should be omitted from index.
     """