feat: add `modify` command

1 files changed, 113 insertions, 15 deletions

diff --git a/app/Main.hs b/app/Main.hs
index 47eeedd..5e48c84 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -12,7 +12,7 @@
 
 module Main where
 
-import Control.Arrow (second)
+import Control.Arrow (first, second, (***))
 import Control.Concurrent.ParallelIO.Local (parallel, withPool)
 import Control.Exception (Exception, throw, throwIO)
 import Control.Monad (forM, unless, when)
@@ -80,6 +80,11 @@ data Cmd
       { indexNames :: [FilePath]
       }
   | TopWords
+  | Modify
+      { indexNames :: [FilePath],
+        tags :: [G.Tag],
+        untags :: [G.Tag]
+      }
 
 args :: O.Parser Args
 args =
@@ -99,7 +104,9 @@ cmd =
       O.command "view" . O.info viewCmd $
         O.progDesc "View document(s)",
       O.command "topwords" . O.info topWordsCmd $
-        O.progDesc "View probability map per tag"
+        O.progDesc "View probability map per tag",
+      O.command "modify" . O.info modifyCmd $
+        O.progDesc "Modify document(s)"
     ]
 
 consumeCmd :: O.Parser Cmd
@@ -138,6 +145,13 @@ topWordsCmd :: O.Parser Cmd
 topWordsCmd =
   pure TopWords
 
+modifyCmd :: O.Parser Cmd
+modifyCmd =
+  Modify
+    <$> indexNamesArg
+    <*> tagsArg
+    <*> untagsArg
+
 inputsArg :: O.Parser [FilePath]
 inputsArg =
   O.many
@@ -185,6 +199,50 @@ filtersArg =
     parse ('!' : '@' : tagKey) = Just (Filter Exclude (FilterByTag (T.pack tagKey)))
     parse _ = Nothing
 
+tagsArg :: O.Parser [G.Tag]
+tagsArg =
+  O.many $
+    O.option
+      (O.maybeReader parse)
+      ( O.long "tag"
+          <> O.help "Tag to add"
+      )
+  where
+    parse ('@' : tag) =
+      let (tagKey, tagValue) =
+            T.strip *** T.strip $
+              T.break (== ' ') (T.pack tag)
+       in Just $
+            G.Tag
+              tagKey
+              ( if T.null tagValue
+                  then Nothing
+                  else (Just tagValue)
+              )
+    parse _ = Nothing
+
+untagsArg :: O.Parser [G.Tag]
+untagsArg =
+  O.many $
+    O.option
+      (O.maybeReader parse)
+      ( O.long "untag"
+          <> O.help "Tag to remove"
+      )
+  where
+    parse ('@' : tag) =
+      let (tagKey, tagValue) =
+            T.strip *** T.strip $
+              T.break (== ' ') (T.pack tag)
+       in Just $
+            G.Tag
+              tagKey
+              ( if T.null tagValue
+                  then Nothing
+                  else (Just tagValue)
+              )
+    parse _ = Nothing
+
 languageArg :: O.Parser (Maybe String)
 languageArg =
   O.optional
@@ -327,6 +385,16 @@ main = do
           )
           M.empty
           (probabilityMap allDocs)
+    Args {cmd = Modify {indexNames, tags, untags}} -> do
+      docs <- mapM (readDocument . (<.> "json")) indexNames
+      mapM_
+        ( \doc -> do
+            let doc' = addTags tags (removeTags untags doc)
+            withGit do
+              J.encodeFile doc'.iFilePath doc'.index
+              commitAll [doc.iFilePath] (printf "tag %s" (takeBaseName doc.iFilePath))
+        )
+        docs
 
 printTags :: Document -> IO ()
 printTags doc =
@@ -567,23 +635,53 @@ processDocuments settings allDocs docs =
         commitAll [doc.iFilePath] (printf "process %s (auto)" doc.iFilePath)
       pure doc'
 
+applyTags :: [Either T.Text G.Tag] -> Document -> Document
+applyTags tags' doc = do
+  addTags tags (removeTags untags doc)
+  where
+    (untags, tags) =
+      first (map (\tagKey -> G.Tag tagKey Nothing))
+        $ partitionEithers tags'
+
+addTags :: [G.Tag] -> Document -> Document
+addTags tags doc =
+  doc
+    { index =
+        doc.index
+          { tags = doc.index.tags `S.union` (S.fromList tags)
+          }
+    }
+
+removeTags :: [G.Tag] -> Document -> Document
+removeTags tags doc =
+  doc
+    { index =
+        doc.index
+          { tags =
+              S.filter
+                ( \tag' ->
+                    ( not $
+                        any
+                          ( \tag ->
+                              maybe
+                                (G.tagKey tag' == G.tagKey tag)
+                                (\_ -> tag' == tag)
+                                (G.tagValue tag)
+                          )
+                          tags
+                    )
+                )
+                doc.index.tags
+          }
+    }
+
 tagDocumentInteractively :: S.Settings -> [Document] -> Document -> IO Document
 tagDocumentInteractively settings allDocs doc = do
   suggestedTags <- suggestTags settings allDocs doc
-  (tagsToRemove, tagsToAdd) <-
-    partitionEithers
-      <$> mapM (uncurry tagDocumentInteractively') suggestedTags
+  tags <- mapM (uncurry tagDocumentInteractively') suggestedTags
   let doc' =
-        doc
-          { index =
-              doc.index
-                { tags =
-                    S.filter
-                      (not . (`elem` tagsToRemove) . G.tagKey)
-                      doc.index.tags
-                      `S.union` (S.fromList tagsToAdd),
-                  todo = False
-                }
+        (applyTags tags doc)
+          { index = doc.index {todo = False}
           }
   withGit do
     J.encodeFile doc.iFilePath doc'.index