added docstrings to LookupDataPool. Added remove_all method to flush …

…LookupDataPool. Added tests for load_data_from_file

flashgeotext/lookup.py

@@ -122,13 +122,34 @@ def validate(self) -> dict:
 
 
 class LookupDataPool:
-    """
+    """Collection of KeywordProcessors from LookupData
+
+    Args:
+        pool (dict): Collection of LookupData.
+
+    Attributes:
+        pool (dict): Collection of LookupData.
+
+    Example:
+        pool = {
+            LookupData.name: flashtext.KeywordProcessor.add_keywords_from_dict(LookupData.data)
+            }
     """
 
     def __init__(self) -> None:
         self.pool: dict = {}
 
     def add(self, lookup: LookupData, update: bool = False) -> None:
+        """Add LookupData to LookupDataPool
+
+        Add LookupData to LookupDataPool.
+        Raises flashgeotext.lookup.LookupDuplicateError if lookup
+        is already in pool unless update == True.
+
+        Args:
+            lookup (LookupData): LookupData to add to pool
+            update (bool): Allow update of an existing entry in LookupDataPool
+        """
         if not isinstance(lookup, LookupData):
             raise TypeError(f"lookup has to be instance of LookupData")
 
@@ -142,11 +163,26 @@ def add(self, lookup: LookupData, update: bool = False) -> None:
             logger.debug(f"{lookup.name} added to pool")
 
     def remove(self, lookup_to_remove: str) -> None:
+        """Remove LookupData from LookupDataPool
+
+        Args:
+            lookup_to_remove (str): LookupData to remove from pool
+        """
         if lookup_to_remove in self.pool:
             del self.pool[lookup_to_remove]
             logger.debug(f"{lookup_to_remove} removed from pool")
 
+    def remove_all(self):
+        """Remove all LookupData from LookupDataPool
+        """
+
+        self.pool = {}
+
     def _add_demo_data(self):
+        """(private) Add demo data to pool
+
+        Adds DEMODATA_CITIES and DEMODATA_COUNTRIES to LookupDataPool
+        """
         cities = LookupData(
             name="cities", data=load_data_from_file(file=DEMODATA_CITIES)
         )
@@ -159,5 +195,16 @@ def _add_demo_data(self):
 
 
 def load_data_from_file(file: str) -> dict:
+    """Load data from json file
+
+    Load data from json file. Raises TypeError if not json
+
+    Args:
+        file (str): path to file
+    """
+
+    if not file.endswith(".json"):
+        raise TypeError("File has to be Filetype .json")
+
     with open(file, "r", encoding="utf-8") as f:
         return json.loads(f.read())

tests/integration/test_geotext_extractor.py

@@ -6,6 +6,13 @@
 text = "Berlin ist die Hauptstadt von Deutschland. Berlin ist nicht haesslich, aber auch nicht sonderlich schoen."
 
 
+def test_geotext_demo_data():
+    geotext = GeoText(use_demo_data=True)
+
+    assert geotext.pool["cities"]
+    assert geotext.pool["countries"]
+
+
 def test_geotext_extract(geotext):
     output = geotext.extract(text)
     assert "Berlin" in output["cities"]

tests/unit/test_lookup_data_pool.py

@@ -1,9 +1,10 @@
 import pytest
 
-from flashgeotext.geotext import GeoText
+from flashgeotext.lookup import load_data_from_file
 from flashgeotext.lookup import LookupData
 from flashgeotext.lookup import LookupDataPool
 from flashgeotext.lookup import LookupDuplicateError
+from flashgeotext.settings import DEMODATA_CITIES
 
 
 def test_lookup_data_pool(test_data_cities):
@@ -56,8 +57,23 @@ def test_lookup_data_pool_remove_lookup_from_pool(test_data_cities):
         assert processor.pool["cities"]
 
 
-def test_lookup_data_pool_with_test_data():
-    geotext = GeoText(use_demo_data=True)
+def test_lookup_data_pool_remove_all_from_pool(test_data_cities):
+    lookup = LookupData(name="cities", data=test_data_cities)
+
+    processor = LookupDataPool()
+    processor.add(lookup)
+
+    processor.remove_all()
 
-    assert geotext.pool["cities"]
-    assert geotext.pool["countries"]
+    assert not processor.pool
+
+
+def test_load_data_from_file():
+    cities = load_data_from_file(DEMODATA_CITIES)
+    assert cities
+
+
+def test_load_data_from_file_raises():
+    with pytest.raises(TypeError):
+        cities = load_data_from_file("demodata_cities.txt")
+        assert cities