Update core.py

Author: wannaphong

pythainlp/corpus/core.py

@@ -88,23 +88,18 @@ def get_corpus(filename: str) -> frozenset:
     return frozenset(lines)
 
 
-def _update_db(name, path):
-    local_db = TinyDB(corpus_db_path())
-    query = Query()
-    local_db.update({"filename": path}, query.name == name)
-    local_db.close()
-
 def _update_all():
     print("Update Corpus...")
-    local_db = TinyDB(corpus_db_path())
-    item_all = local_db.all()
+    with TinyDB(corpus_db_path()) as local_db:
+        item_all = local_db.all()
+        query = Query()
+        for item in item_all:
+            name = item['name']
+            if 'file_name' in item.keys():
+                local_db.update({"filename": item['file_name']}, query.name == name)
+            elif 'file' in item.keys():
+                local_db.update({"filename": item['file']}, query.name == name)
     local_db.close()
-    for item in item_all:
-        name = item['name']
-        if 'file_name' in item.keys():
-            _update_db(name, item['file_name'])
-        elif 'file' in item.keys():
-            _update_db(name, item['file'])
 
 def get_corpus_path(name: str) -> Union[str, None]:
     """
@@ -143,9 +138,10 @@ def get_corpus_path(name: str) -> Union[str, None]:
     """
     # check if the corpus is in local catalog, download if not
     corpus_db_detail = get_corpus_db_detail(name)
-    if (corpus_db_detail.get("file_name") is not None or corpus_db_detail.get("file") is not None) and corpus_db_detail.get("filename") is None:
+    if corpus_db_detail.get("file_name") is not None and corpus_db_detail.get("filename") is None:
+        _update_all()
+    elif corpus_db_detail.get("file") is not None and corpus_db_detail.get("filename") is None:
         _update_all()
-
 
     if not corpus_db_detail or not corpus_db_detail.get("filename"):
         download(name)