From bdcf5aa36607286699fd397d3cc9c39f9d322086 Mon Sep 17 00:00:00 2001
From: voronin9032 <voronin9032n3@gmail.com>
Date: Sat, 18 Jun 2022 05:56:58 +0300
Subject: [PATCH] Add logging to CLI

---
 pyproject.toml      |  2 +-
 traps/__main__.py   |  4 ++--
 traps/cli.py        | 12 +++++++++---
 traps/downloader.py | 17 ++++++++++++++++-
 traps/utils.py      |  4 ++++
 5 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index a485adf..1b6da9d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ loguru = "^0.6.0"
 [tool.poetry.dev-dependencies]
 
 [tool.poetry.scripts]
-traps = "traps.cli:cli"
+traps = "traps.cli:main"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
diff --git a/traps/__main__.py b/traps/__main__.py
index 98dcca0..9ae637f 100644
--- a/traps/__main__.py
+++ b/traps/__main__.py
@@ -1,4 +1,4 @@
-from .cli import cli
+from .cli import main
 
 if __name__ == "__main__":
-    cli()
+    main()
diff --git a/traps/cli.py b/traps/cli.py
index 52c09ac..8ba3caa 100644
--- a/traps/cli.py
+++ b/traps/cli.py
@@ -26,6 +26,7 @@ def cli(verbose: bool):
     else:
         loglevel = "INFO"
 
+    # Remove the default logger if it exists.
     try:
         logger.remove(0)
     except ValueError:
@@ -35,7 +36,8 @@ def cli(verbose: bool):
         sys.stderr,
         level=loglevel,
         format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>"
-               "{level: <8}</level> | <level>{message}</level>"
+               "{level: <8}</level> | <level>{message}</level>",
+        filter=lambda record: record["extra"].get("name") == "traps-logger"
     )
 
 
@@ -50,9 +52,13 @@ def install(directory: pathlib.Path, amount: int):
 
 @cli.command("version", help="Print version and exit.")
 def version():
-    print(f"traps {traps.__version__}")
+    click.echo(f"{traps.__name__} {traps.__version__}")
     sys.exit(0)
 
 
+def main():
+    cli.main(windows_expand_args=False)
+
+
 if __name__ == "__main__":
-    cli()
+    main()
diff --git a/traps/downloader.py b/traps/downloader.py
index c2a3c24..ca2d71f 100644
--- a/traps/downloader.py
+++ b/traps/downloader.py
@@ -8,6 +8,7 @@ import requests
 from click import BadParameter
 
 from traps.utils import filename_from_url
+from traps.utils import logger
 
 __all__ = ["get"]
 API_URL = "https://safebooru.org/index.php"
@@ -19,6 +20,7 @@ def _fetch_urls(n: int = 1) -> List[str]:
         raise BadParameter("you can't download more than 5000 files at a time")
     if n < 1:
         raise BadParameter("you can't download a negative number of files")
+
     used_offsets = []
     urls = []
 
@@ -43,6 +45,7 @@ def _fetch_urls(n: int = 1) -> List[str]:
             for _ in range(limit)
         ]
 
+    logger.info("Fetching urls")
     if n > 100:
         with ThreadPoolExecutor(max_workers=16) as p:
             for i in p.map(lambda _: fetch(100), range(n // 100)):
@@ -50,23 +53,35 @@ def _fetch_urls(n: int = 1) -> List[str]:
             n %= 100
     if n < 100:
         urls += fetch(n)
+    logger.info("Done")
     return urls
 
 
 def _download(directory: Path, url: str) -> None:
+    logger.debug(f"Downloading {url}")
     resp = requests.get(url, stream=True)
+    if not resp.ok:
+        logger.warning(f"Couldn't download {url}: HTTP error {resp.status_code}")
+        return
     filename = filename_from_url(url)
     with open(directory / filename, "wb") as f:
         for part in resp.iter_content(1024):
             if not part:
                 break
             f.write(part)
+        else:
+            logger.success(f"Downloaded {url}")
 
 
 def get(directory: Union[str, Path] = "traps", amount: int = 1) -> None:
     if not isinstance(directory, Path):
         directory = Path(directory)
-    directory.mkdir(exist_ok=True)
+    if not directory.exists():
+        logger.debug(f"Creating {directory.name} directory")
+        directory.mkdir()
+        logger.debug("Done")
     urls = _fetch_urls(amount)
+    logger.info("Downloading traps")
     with ThreadPoolExecutor(max_workers=16) as p:
         p.map(lambda url: _download(directory, url), urls)
+    logger.success(f"Downloaded {amount} traps")
diff --git a/traps/utils.py b/traps/utils.py
index 3f60979..53fa0d9 100644
--- a/traps/utils.py
+++ b/traps/utils.py
@@ -1,6 +1,10 @@
 import pathlib
 import urllib.parse
 
+import loguru
+
+logger = loguru.logger.bind(name="traps-logger")
+
 
 def filename_from_url(url: str) -> str:
     path = urllib.parse.urlparse(url).path