Merge pull request #657 from PyThaiNLP/dev

Update from dev

Author: wannaphong

pythainlp/corpus/core.py

@@ -14,6 +14,9 @@
 from pythainlp.tools import get_full_data_path
 from requests.exceptions import HTTPError
 from tinydb import Query, TinyDB
+import tarfile
+import zipfile
+import shutil
 from pythainlp import __version__
 
 
@@ -207,7 +210,10 @@ def get_corpus_path(name: str,  version: str = None) -> Union[str, None]:
 
     if corpus_db_detail and corpus_db_detail.get("filename"):
         # corpus is in the local catalog, get full path to the file
-        path = get_full_data_path(corpus_db_detail.get("filename"))
+        if corpus_db_detail.get("is_folder"):
+            path = get_full_data_path(corpus_db_detail.get("foldername"))
+        else:
+            path = get_full_data_path(corpus_db_detail.get("filename"))
         # check if the corpus file actually exists, download if not
         if not os.path.exists(path):
             download(name)
@@ -411,11 +417,45 @@ def download(
                 file_name, corpus_versions["md5"],
             )
 
+            is_folder = False
+            foldername = None
+
+            if corpus_versions["is_tar_gz"] == "True":
+                is_folder = True
+                foldername = name+"_"+str(version)
+                if not os.path.exists(get_full_data_path(foldername)):
+                    os.mkdir(get_full_data_path(foldername))
+                with tarfile.open(get_full_data_path(file_name)) as tar:
+                    tar.extractall(path=get_full_data_path(foldername))
+            elif corpus_versions["is_zip"] == "True":
+                is_folder = True
+                foldername = name+"_"+str(version)
+                if not os.path.exists(get_full_data_path(foldername)):
+                    os.mkdir(get_full_data_path(foldername))
+                with zipfile.ZipFile(
+                    get_full_data_path(file_name), 'r'
+                ) as zip:
+                    zip.extractall(path=get_full_data_path(foldername))
+
             if found:
-                local_db.update({"version": version}, query.name == name)
+                local_db.update(
+                    {
+                        "version": version,
+                        "filename": file_name,
+                        "is_folder": is_folder,
+                        "foldername": foldername
+                    },
+                    query.name == name
+                )
             else:
                 local_db.insert(
-                    {"name": name, "version": version, "filename": file_name}
+                    {
+                        "name": name,
+                        "version": version,
+                        "filename": file_name,
+                        "is_folder": is_folder,
+                        "foldername": foldername
+                    }
                 )
         else:
             if local_db.search(
@@ -471,7 +511,11 @@ def remove(name: str) -> bool:
 
     if data:
         path = get_corpus_path(name)
-        os.remove(path)
+        if data[0].get("is_folder"):
+            os.remove(get_full_data_path(data[0].get("filename")))
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            os.remove(path)
         db.remove(query.name == name)
         db.close()
         return True

pythainlp/translate/en_th.py

@@ -27,19 +27,12 @@
 
 
 def _get_translate_path(model: str, *path: str) -> str:
-    return os.path.join(get_full_data_path(model), *path)
+    return os.path.join(get_corpus_path(model, version="1.0"), *path)
 
 
 def _download_install(name: str) -> None:
     if get_corpus_path(name) is None:
         download(name, force=True, version="1.0")
-        tar = tarfile.open(get_corpus_path(name, version="1.0"), "r:gz")
-        tar.extractall()
-        tar.close()
-    if not os.path.exists(get_full_data_path(name)):
-        os.mkdir(get_full_data_path(name))
-        with tarfile.open(get_corpus_path(name)) as tar:
-            tar.extractall(path=get_full_data_path(name))
 
 
 def download_model_all() -> None:

tests/test_corpus.py

@@ -28,6 +28,7 @@
 from pythainlp.corpus.util import revise_newmm_default_wordset
 from requests import Response
 import nltk
+import os
 
 
 class TestCorpusPackage(unittest.TestCase):
@@ -161,3 +162,8 @@ def test_revise_wordset(self):
             ["ที่", "ถูก", "สังหาร", "เมื่อ", "ปี", " ", "พ.ศ.", " ", "2492"],
         ]
         self.assertIsInstance(revise_newmm_default_wordset(training_data), set)
+
+    def test_zip(self):
+        _p = get_corpus_path("test_zip")
+        self.assertEqual(os.path.isdir(_p), True)
+        self.assertEqual(remove("test_zip"), True)