From 1c385e1b2a42f04a6475c44a8b9c0af3bf2a69d7 Mon Sep 17 00:00:00 2001
From: vzhd1701 <vzhd1701@gmail.com>
Date: Thu, 7 Apr 2022 11:29:13 +0500
Subject: [PATCH] feat: add option to set custom tag for uploaded pages

fix #23
---
 enex2notion/cli.py | 10 ++++++++++
 tests/test_cli.py  |  6 ++++++
 2 files changed, 16 insertions(+)

diff --git a/enex2notion/cli.py b/enex2notion/cli.py
index a2b6ccf..da5df7c 100644
--- a/enex2notion/cli.py
+++ b/enex2notion/cli.py
@@ -64,6 +64,7 @@ def __init__(
         add_meta: bool,
         add_pdf_preview: bool,
         condense_lines: bool,
+        custom_tag: str,
     ):
         self.import_root = import_root
         self.mode = mode
@@ -72,6 +73,7 @@ def __init__(
         self.add_meta = add_meta
         self.add_pdf_preview = add_pdf_preview
         self.condense_lines = condense_lines
+        self.custom_tag = custom_tag
 
     def upload(self, enex_file: Path):
         logger.info(f"Processing notebook '{enex_file.stem}'...")
@@ -83,6 +85,9 @@ def upload(self, enex_file: Path):
                 logger.debug(f"Skipping note '{note.title}' (already uploaded)")
                 continue
 
+            if self.custom_tag and self.custom_tag not in note.tags:
+                note.tags.append(self.custom_tag)
+
             note_blocks = self._parse_note(note)
             if not note_blocks:
                 continue
@@ -132,6 +137,7 @@ def cli(argv):
         add_meta=args.add_meta,
         add_pdf_preview=args.add_pdf_preview,
         condense_lines=args.condense_lines,
+        custom_tag=args.tag,
     )
 
     for enex_input in args.enex_input:
@@ -227,6 +233,10 @@ def parse_args(argv):
                 " makes sense only with PAGE mode"
             ),
         },
+        "--tag": {
+            "type": str,
+            "help": "add custom tag to uploaded notes",
+        },
         "--condense-lines": {
             "action": "store_true",
             "default": False,
diff --git a/tests/test_cli.py b/tests/test_cli.py
index f99693f..df9675d 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -257,5 +257,11 @@ def test_bad_token(mock_api, fake_note_factory, caplog):
     assert "Invalid token provided!" in caplog.text
 
 
+def test_custom_tag(mock_api, fake_note_factory):
+    cli(["--tag", "test_tag", "fake.enex"])
+
+    fake_note_factory()[0].tags.append.assert_called_once_with("test_tag")
+
+
 def test_cli_main_import():
     from enex2notion import __main__