Add options to CLI

poetry.lock

@@ -17,6 +17,26 @@ python-versions = ">=3.5.0"
 [package.extras]
 unicode_backport = ["unicodedata2"]
 
+[[package]]
+name = "click"
+version = "8.1.3"
+description = "Composable command line interface toolkit"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+
+[package.dependencies]
+colorama = {version = "*", markers = "platform_system == \"Windows\""}
+importlib-metadata = {version = "*", markers = "python_version < \"3.8\""}
+
+[[package]]
+name = "colorama"
+version = "0.4.4"
+description = "Cross-platform colored terminal text."
+category = "main"
+optional = false
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+
 [[package]]
 name = "idna"
 version = "3.3"
@@ -25,6 +45,23 @@ category = "main"
 optional = false
 python-versions = ">=3.5"
 
+[[package]]
+name = "importlib-metadata"
+version = "4.11.4"
+description = "Read metadata from Python packages"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+
+[package.dependencies]
+typing-extensions = {version = ">=3.6.4", markers = "python_version < \"3.8\""}
+zipp = ">=0.5"
+
+[package.extras]
+docs = ["sphinx", "jaraco.packaging (>=9)", "rst.linker (>=1.9)"]
+perf = ["ipython"]
+testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-flake8", "pytest-cov", "pytest-enabler (>=1.0.1)", "packaging", "pyfakefs", "flufl.flake8", "pytest-perf (>=0.9.2)", "pytest-black (>=0.3.7)", "pytest-mypy (>=0.9.1)", "importlib-resources (>=1.3)"]
+
 [[package]]
 name = "requests"
 version = "2.28.0"
@@ -43,6 +80,14 @@ urllib3 = ">=1.21.1,<1.27"
 socks = ["PySocks (>=1.5.6,!=1.5.7)"]
 use_chardet_on_py3 = ["chardet (>=3.0.2,<5)"]
 
+[[package]]
+name = "typing-extensions"
+version = "4.2.0"
+description = "Backported and Experimental Type Hints for Python 3.7+"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+
 [[package]]
 name = "urllib3"
 version = "1.26.9"
@@ -56,10 +101,22 @@ brotli = ["brotlicffi (>=0.8.0)", "brotli (>=1.0.9)", "brotlipy (>=0.6.0)"]
 secure = ["pyOpenSSL (>=0.14)", "cryptography (>=1.3.4)", "idna (>=2.0.0)", "certifi", "ipaddress"]
 socks = ["PySocks (>=1.5.6,!=1.5.7,<2.0)"]
 
+[[package]]
+name = "zipp"
+version = "3.8.0"
+description = "Backport of pathlib-compatible object wrapper for zip files"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+
+[package.extras]
+docs = ["sphinx", "jaraco.packaging (>=9)", "rst.linker (>=1.9)"]
+testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-flake8", "pytest-cov", "pytest-enabler (>=1.0.1)", "jaraco.itertools", "func-timeout", "pytest-black (>=0.3.7)", "pytest-mypy (>=0.9.1)"]
+
 [metadata]
 lock-version = "1.1"
 python-versions = "^3.7"
-content-hash = "bfd577599bdf37e40d0c9ec345e47b18bbb479ca31b547d8292acf1dbf7d9727"
+content-hash = "1d7efb30a6917067ee323aff51674a885542ec0ab42d54f1b3a07ea5e0aac6e2"
 
 [metadata.files]
 certifi = [
@@ -70,15 +127,35 @@ charset-normalizer = [
     {file = "charset-normalizer-2.0.12.tar.gz", hash = "sha256:2857e29ff0d34db842cd7ca3230549d1a697f96ee6d3fb071cfa6c7393832597"},
     {file = "charset_normalizer-2.0.12-py3-none-any.whl", hash = "sha256:6881edbebdb17b39b4eaaa821b438bf6eddffb4468cf344f09f89def34a8b1df"},
 ]
+click = [
+    {file = "click-8.1.3-py3-none-any.whl", hash = "sha256:bb4d8133cb15a609f44e8213d9b391b0809795062913b383c62be0ee95b1db48"},
+    {file = "click-8.1.3.tar.gz", hash = "sha256:7682dc8afb30297001674575ea00d1814d808d6a36af415a82bd481d37ba7b8e"},
+]
+colorama = [
+    {file = "colorama-0.4.4-py2.py3-none-any.whl", hash = "sha256:9f47eda37229f68eee03b24b9748937c7dc3868f906e8ba69fbcbdd3bc5dc3e2"},
+    {file = "colorama-0.4.4.tar.gz", hash = "sha256:5941b2b48a20143d2267e95b1c2a7603ce057ee39fd88e7329b0c292aa16869b"},
+]
 idna = [
     {file = "idna-3.3-py3-none-any.whl", hash = "sha256:84d9dd047ffa80596e0f246e2eab0b391788b0503584e8945f2368256d2735ff"},
     {file = "idna-3.3.tar.gz", hash = "sha256:9d643ff0a55b762d5cdb124b8eaa99c66322e2157b69160bc32796e824360e6d"},
 ]
+importlib-metadata = [
+    {file = "importlib_metadata-4.11.4-py3-none-any.whl", hash = "sha256:c58c8eb8a762858f49e18436ff552e83914778e50e9d2f1660535ffb364552ec"},
+    {file = "importlib_metadata-4.11.4.tar.gz", hash = "sha256:5d26852efe48c0a32b0509ffbc583fda1a2266545a78d104a6f4aff3db17d700"},
+]
 requests = [
     {file = "requests-2.28.0-py3-none-any.whl", hash = "sha256:bc7861137fbce630f17b03d3ad02ad0bf978c844f3536d0edda6499dafce2b6f"},
     {file = "requests-2.28.0.tar.gz", hash = "sha256:d568723a7ebd25875d8d1eaf5dfa068cd2fc8194b2e483d7b1f7c81918dbec6b"},
 ]
+typing-extensions = [
+    {file = "typing_extensions-4.2.0-py3-none-any.whl", hash = "sha256:6657594ee297170d19f67d55c05852a874e7eb634f4f753dbd667855e07c1708"},
+    {file = "typing_extensions-4.2.0.tar.gz", hash = "sha256:f1c24655a0da0d1b67f07e17a5e6b2a105894e6824b92096378bb3668ef02376"},
+]
 urllib3 = [
     {file = "urllib3-1.26.9-py2.py3-none-any.whl", hash = "sha256:44ece4d53fb1706f667c9bd1c648f5469a2ec925fcf3a776667042d645472c14"},
     {file = "urllib3-1.26.9.tar.gz", hash = "sha256:aabaf16477806a5e1dd19aa41f8c2b7950dd3c746362d7e3223dbe6de6ac448e"},
 ]
+zipp = [
+    {file = "zipp-3.8.0-py3-none-any.whl", hash = "sha256:c4f6e5bbf48e74f7a38e7cc5b0480ff42b0ae5178957d564d18932525d5cf099"},
+    {file = "zipp-3.8.0.tar.gz", hash = "sha256:56bf8aadb83c24db6c4b577e13de374ccfb67da2078beba1d037c17980bf43ad"},
+]

pyproject.toml

@@ -23,6 +23,7 @@ classifiers = [
 [tool.poetry.dependencies]
 python = "^3.7"
 requests = "^2.28.0"
+click = "^8.1.3"
 
 [tool.poetry.dev-dependencies]
 

traps/__init__.py

@@ -1,28 +1,19 @@
-import requests
 import pathlib
-from random import randint
+import secrets
 import urllib.parse
+from multiprocessing import Pool
 from threading import Thread
 
-api_url = "https://api.waifu.pics/nsfw/trap"
+import click
+import requests
+
+API_URL = "https://api.waifu.pics/nsfw/trap"
 
 
-def get(directory="traps", amount=randint(5, 10)):
-    path = pathlib.Path(directory)
-    if not path.exists():
-        path.mkdir()
-    for i in range(amount):
-        Thread(target=traps, args=(directory,)).start()
-
-
-def traps(directory):
-    req_url = requests.get(api_url)
-    url = req_url.json()["url"]
-    if not req_url.ok:
-        print("error:", req_url)
-        return
+def get_traps(params):
+    directory, url = params
     filename = urllib.parse.urlparse(url)
-    filename = pathlib.Path(directory, pathlib.Path(filename.path).name)
+    filename = directory.joinpath(secrets.token_hex(8) + pathlib.Path(filename.path).suffix)
     with open(filename, "wb") as f:
         response = requests.get(url, stream=True)
         for block in response.iter_content(1024):
@@ -31,8 +22,48 @@ def traps(directory):
             f.write(block)
 
 
-def main():
-    get()
+def fetch_url(urls_list: list):
+    url = requests.get(API_URL).json()["url"]
+    urls_list.append(url)
+
+
+@click.command(help="how about you pip install some traps")
+@click.option(
+    "-n",
+    "-t",
+    "--traps",
+    type=click.INT,
+    default=10,
+    show_default=True,
+    help="number of traps to get"
+)
+@click.argument(
+    "directory",
+    default="traps",
+    type=click.Path(
+        dir_okay=True,
+        file_okay=False,
+        path_type=pathlib.Path
+    )
+)
+def main(traps: int, directory: pathlib.Path):
+    if not directory.exists():
+        directory.mkdir()
+    urls = []
+    threads = [
+        Thread(target=fetch_url, args=(urls,))
+        for _ in range(traps)
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+    params = [
+        (directory, url)
+        for url in urls
+    ]
+    with Pool(8) as p:
+        p.map(get_traps, params)
 
 
 if __name__ == '__main__':