Fix variable naming

README.md

@@ -20,40 +20,17 @@ active = extract
 
 [[extract]]
 on_import = False
-formatting = "[{page}]{newline}{quote_begin}> {quote}{newline}{quote_end}{note_begin}Note: {note}{note_end}"
-minimum_similarity = 0.75
+minimum_text_similarity = 0.75
+minimum_color_similarity = 0.833
+formatting = "{%quote_container> {quote} %}[{page}]{%note_container{newline}Note: {note} %}{%tag_container #{tag}%}"
 ```
 
 If `on_import` is `True` extraction is automatically run whenever a new document is added to the library,
 if false extraction has to be handled manually.
 
-`formatting` takes a string with a variety of template options. You can use any of the following:
+---
 
-- `{page}`: The page number the annotation was found on.
-- `{quote}`: The actual quoted string (i.e. highlighted).
-- `{note}`: The annotation note (i.e. addded reader).
-- `{quote_begin}`: Mark the area that begins a quotation. Useful to get rid of prefixes or suffixes if no quotation exists. Requires a corresponding `{quote_end}` to be set after.
-- `{note_begin}`: Same idea for the note area. Requires a corresponding `{note_end}` to be set after.
-
-For example, the default formatting string will result in this output:
-
-```markdown
-[4]
-> came “urban rights” caused collective outrage. Thus, through- out 2007 and 2008, protestors in towns and villages across the Nile Delta poured into the streets to rally over cuts in water flow. Deployment of massive riot police could not stop
-Note: Often illegally connected to network, ‘revolution of the thirsty’
-```
-
-The formatting string `{quote_begin}Quote: {quote}{page}{newline}{quote_end}{note_begin}Note: {note}{note_end}`
-will result in the following output:
-
-```markdown
-Quote: came “urban rights” caused collective outrage. Thus, through- out 2007 and 2008, protestors in towns and villages across the Nile Delta poured into the streets to rally over cuts in water flow. Deployment of massive riot police could not stop [4]
-Note: Often illegally connected to network, ‘revolution of the thirsty’
-```
-
-Note, however, that in this example the page number if *only* displayed if the annotation contains a quote.
-
-`minimum_similarity` sets the required similarity of an annotation's note and written words to be viewed
+`minimum_text_similarity` sets the required similarity of an annotation's note and written words to be viewed
 as one. Any annotation that has both and is *under* the minimum similarity will be added in the following form:
 
 ```markdown
@@ -64,6 +41,47 @@ Note: my additional thoughts
 That is, the extractor detects additional written words by whoever annotated and adds them to the extraction.
 The option generally should not take too much tuning, but it is there if you need it.
 
+---
+
+`minimum_color_similarity` sets the required similarity of highlight/annotation colors to be recognized as the 'pure' versions of themselves for color mapping (see below). With a low required similarity, for example dark green and light green will both be recognized simply as 'green' while a high similarity will not match them, instead only matching closer matches to a pure (0, 255, 0) green value.
+
+This should generally be an alright default but is here to be changed for example if you work with a lot of different annotation colors (where dark purple and light purple get different meanings) and get false positives.
+
+---
+
+The plugin contains a configuration sub-category of `tags`: Here, you can define meaning for your highlight/annotation colors. For example, if you always highlight the main arguments and findings in orange and always highlight things you have to follow up on in blue, you can assign the meanings 'important' and 'todo' to them respectively as follows:
+
+```ini
+[[[tags]]]
+orange = "important"
+blue = "todo"
+```
+
+Currently recognized colors are: `red` `green` `blue` `yellow` `purple` `orange`.
+Since these meanings are often highly dependent on personal organization and reading systems,
+no defaults are set here.
+
+---
+
+`formatting` takes a string with a variety of template options. You can use any of the following:
+
+- `{page}`: The page number the annotation was found on.
+- `{quote}`: The actual quoted string (i.e. highlighted).
+- `{note}`: The annotation note (i.e. addded reader).
+- `{%quote_container [other text] %}`: Mark the area that contains a quotation. Useful to get rid of prefixes or suffixes if no quotation exists. Usually contains some plain text and a `{quote}` element. Can *not* be nested with other containers.
+- `{%note_container [other text] %}`: Mark the area that contains a note. Useful to get rid of prefixes or suffixes if no note exists. Usually contains some plain text and a `{note}` element. Can *not* be nested with other containers.
+- `{%tag_container [other text] %}`: Mark the area that contains a tag. Useful to get rid of prefixes or suffixes if no tag exists. Usually contains some plain text and a `{tag}` element. Can *not* be nested with other containers.
+- `{newline}`: Add a line break in the resulting annotation display.
+
+For example, the default formatting string `"{%quote_container> {quote} %}[{page}]{%note_container{newline}Note: {note} %}{%tag_container #{tag}%}"` will result in this output:
+
+```
+> Mobilizing the TPSN scheme (see above) and drawing on cultural political economy and critical governance studies, this landmark article offers an alternative account [5]
+Note: a really intersting take on polydolywhopp
+```
+
+Container marks are useful to encapsulate a specific type of the annotation, so extracted annotations in the end don't contains useless linebreaks or quotations markup.
+
 ## Usage:
 
 `pubs extract [-h|-w|-e] <citekeys>`
@@ -131,20 +149,21 @@ content, because then we can just use that. It is harder to parse if it does not
 ## Roadmap:
 
 - [x] extracts highlights and annotations from a doc file (e.g. using PyMuPDF)
-- [ ] puts those in the annotation file of a doc in a customizable format
+- [x] puts those in the annotation file of a doc in a customizable format
 - [x] option to have it automatically run after a file is added?
     - option to have it run whenever a pdf in the library was updated?
 - [ ] needs some way to delimit where it puts stuff and user stuff is in note
     - [ ] one way is to have it look at `> [17] here be extracted annotation from page seventeen` annotations and put it in between
     - [x] another, probably simpler first, is to just append missing annotations to the end of the note
-- [ ] some highlights (or annotations in general) do not contain text as content
-    - [ ] pymupdf can extract the content of the underlying rectangle (mostly)
-    - [ ] issue is that sometimes the highlight contents are in content, sometimes a user comment instead
-        - [ ] we could have a comparison function which estimates how 'close' the two text snippets are and act accordingly
-- [ ] config option to map colors in annotations to meaning ('read', 'important', 'extra') in pubs
-    - [ ] colors are given in very exact 0.6509979 RGB values, meaning we could once again estimate if a color is 'close enough' in distance to tag it accordingly
-- [ ] make invoking the command run a query if corresponding option provided (or whatever) in pubs syntax and use resulting papers
-    - [ ] confirm for many papers?
+- [x] some highlights (or annotations in general) do not contain text as content
+    - [x] pymupdf can extract the content of the underlying rectangle (mostly)
+    - [x] issue is that sometimes the highlight contents are in content, sometimes a user comment instead
+        - [x] we could have a comparison function which estimates how 'close' the two text snippets are and act accordingly -> using levenshtein distance
+- [x] config option to map colors in annotations to meaning ('read', 'important', 'extra') in pubs
+    - [x] colors are given in very exact 0.6509979 RGB values, meaning we could once again estimate if a color is 'close enough' in distance to tag it accordingly -> using euclidian distance
+    - [ ] support custom colors by setting a float tuple in configuration
+- [x] make invoking the command run a query if corresponding option provided (or whatever) in pubs syntax and use resulting papers
+    - [x] confirm for many papers?
 - [ ] warning when the amount of annotations in file is different than the amount extracted?
 
 ## Things that would also be nice in pubs in general and don't really belong in this repository

extract/extract.py

@@ -1,11 +1,15 @@
 import os
 import re
 import argparse
+import math
+from dataclasses import dataclass, field
+from typing import Dict
 
 import fitz
 import Levenshtein
 
 from pubs.plugins import PapersPlugin
+from pubs.paper import Paper
 from pubs.events import DocAddEvent, NoteEvent
 
 from pubs import repo, pretty
@@ -13,7 +17,89 @@ from pubs.utils import resolve_citekey_list
 from pubs.content import check_file, read_text_file, write_file
 from pubs.query import get_paper_filter
 
-CONFIRMATION_PAPER_THRESHOLD=5
+CONFIRMATION_PAPER_THRESHOLD = 5
+TEXT_SIMILARITY_MINIMUM = 0.75
+COLOR_SIMILARITY_MINIMUM = 0.833
+
+COLORS = {
+    "red": (1, 0, 0),
+    "green": (0, 1, 0),
+    "blue": (0, 0, 1),
+    "yellow": (1, 1, 0),
+    "purple": (0.5, 0, 0.5),
+    "orange": (1, 0.65, 0),
+}
+
+
+@dataclass
+class Annotation:
+    """A PDF annotation object"""
+
+    paper: Paper
+    file: str
+    type: str = "Highlight"
+    text: str = ""
+    content: str = ""
+    page: int = 1
+    colors: Dict = field(default_factory=lambda: {"stroke": (0.0, 0.0, 0.0)})
+    tag: str = ""
+
+    def format(self, formatting):
+        """Return a formatted string of the annotation.
+
+        Given a provided formatting pattern, this method returns the annotation
+        formatted with the correct marker replacements and removals, ready
+        for display or writing.
+        """
+        output = formatting
+        replacements = {
+            r"{quote}": self.text,
+            r"{note}": self.content,
+            r"{page}": str(self.page),
+            r"{newline}": "\n",
+            r"{tag}": self.tag,
+        }
+        pattern = re.compile(
+            "|".join(
+                [re.escape(k) for k in sorted(replacements, key=len, reverse=True)]
+            ),
+            flags=re.DOTALL,
+        )
+        patt_quote_container = re.compile(r"{%quote_container(.*?)%}")
+        patt_note_container = re.compile(r"{%note_container(.*?)%}")
+        patt_tag_container = re.compile(r"{%tag_container(.*?)%}")
+        output = patt_quote_container.sub(r"\1" if self.text else "", output)
+        output = patt_note_container.sub(r"\1" if self.content else "", output)
+        output = patt_tag_container.sub(r"\1" if self.tag else "", output)
+        return pattern.sub(lambda x: replacements[x.group(0)], output)
+
+    @property
+    def colorname(self):
+        """Return the stringified version of the annotation color.
+
+        Finds the closest named color to the annotation and returns it.
+        """
+        annot_colors = (
+            self.colors.get("stroke") or self.colors.get("fill") or (0.0, 0.0, 0.0)
+        )
+        nearest = None
+        minimum_similarity = COLOR_SIMILARITY_MINIMUM
+        for name, values in COLORS.items():
+            similarity_ratio = self._color_similarity_ratio(values, annot_colors)
+            if similarity_ratio > minimum_similarity:
+                minimum_similarity = similarity_ratio
+                nearest = name
+        return nearest
+
+    def _color_similarity_ratio(self, color_one, color_two):
+        """Return the similarity of two colors between 0 and 1.
+
+        Takes two rgb color tuples made of floats between 0 and 1, e.g. (1, 0.65, 0) for orange,
+        and returns the similarity between them, with 1 being the same color and 0 being the
+        difference between full black and full white, as a float.
+        """
+        return 1 - (abs(math.dist([*color_one], [*color_two])) / 3)
+
 
 class ExtractPlugin(PapersPlugin):
     """Extract annotations from any pdf document.
@@ -40,11 +126,17 @@ class ExtractPlugin(PapersPlugin):
 
         settings = conf["plugins"].get("extract", {})
         self.on_import = settings.get("on_import", False)
-        self.minimum_similarity = float(settings.get("minimum_similarity", 0.75))
+        self.minimum_similarity = float(
+            settings.get("minimum_text_similarity", TEXT_SIMILARITY_MINIMUM)
+        )
+        self.minimum_color_similarity = float(
+            settings.get("minimum_color_similarity", COLOR_SIMILARITY_MINIMUM)
+        )
         self.formatting = settings.get(
             "formatting",
-            "{newline}{quote_begin}> {quote} {quote_end}[{page}]{note_begin}{newline}Note: {note}{note_end}",
+            "{%quote_container> {quote} %}[{page}]{%note_container{newline}Note: {note} %}{%tag_container #{tag}%}",
         )
+        self.color_mapping = settings.get("tags", {})
 
     def update_parser(self, subparsers, _):
         """Allow the usage of the pubs extract subcommand"""
@@ -116,15 +208,19 @@ class ExtractPlugin(PapersPlugin):
         Returns all annotations belonging to the papers that
         are described by the citekeys passed in.
         """
-        papers_annotated = []
+        papers_annotated = {}
         for paper in papers:
             file = self._get_file(paper)
             try:
-                papers_annotated.append((paper, self._get_annotations(file)))
+                annotations = self._get_annotations(file, paper)
+                papers_annotated[paper.citekey] = annotations
             except fitz.FileDataError as e:
                 self.ui.error(f"Document {file} is broken: {e}")
         return papers_annotated
 
+    def tag_from_colorname(self, colorname):
+        return self.color_mapping.get(colorname, "")
+
     def _gather_papers(self, conf, args):
         """Get all papers for citekeys.
 
@@ -177,7 +273,7 @@ class ExtractPlugin(PapersPlugin):
             self.ui.warning(f"{paper.citekey} has no valid document.")
         return path
 
-    def _get_annotations(self, filename):
+    def _get_annotations(self, filename, paper):
         """Extract annotations from a file.
 
         Returns all readable annotations contained in the file
@@ -189,35 +285,19 @@ class ExtractPlugin(PapersPlugin):
             for page in doc:
                 for annot in page.annots():
                     quote, note = self._retrieve_annotation_content(page, annot)
-                    annotations.append(
-                        self._format_annotation(quote, note, page.number or 0)
+                    a = Annotation(
+                        file=filename,
+                        paper=paper,
+                        text=quote,
+                        content=note,
+                        colors=annot.colors,
+                        type=annot.type[1],
+                        page=(page.number or 0) + 1,
                     )
+                    a.tag = self.tag_from_colorname(a.colorname)
+                    annotations.append(a)
         return annotations
 
-    def _format_annotation(self, quote, note, pagenumber=0):
-        output = self.formatting
-        replacements = {
-            r"{quote}": quote,
-            r"{note}": note,
-            r"{page}": str(pagenumber),
-            r"{newline}": "\n",
-        }
-        if note == "":
-            output = re.sub(r"{note_begin}.*{note_end}", "", output)
-        if quote == "":
-            output = re.sub(r"{quote_begin}.*{quote_end}", "", output)
-        output = re.sub(r"{note_begin}", "", output)
-        output = re.sub(r"{note_end}", "", output)
-        output = re.sub(r"{quote_begin}", "", output)
-        output = re.sub(r"{quote_end}", "", output)
-        pattern = re.compile(
-            "|".join(
-                [re.escape(k) for k in sorted(replacements, key=len, reverse=True)]
-            ),
-            flags=re.DOTALL,
-        )
-        return pattern.sub(lambda x: replacements[x.group(0)], output)
-
     def _retrieve_annotation_content(self, page, annotation):
         """Gets the text content of an annotation.
 
@@ -249,15 +329,13 @@ class ExtractPlugin(PapersPlugin):
         ready to be passed on through pipelines etc.
         """
         output = ""
-        for contents in annotated_papers:
-            paper = contents[0]
-            annotations = contents[1]
-            if annotations:
-                output += f"------ {paper.citekey} ------\n"
-                for annot in annotations:
-                    output += f"{annot}\n"
+        for citekey, annotations in annotated_papers.items():
+            output += f"------ {citekey} ------\n"
+            for annotation in annotations:
+                # for annot in annotations:
+                output += f"{annotation.format(self.formatting)}\n"
                 output += "\n"
-        print(output)
+        self.ui.message(output)
 
     def _to_notes(self, annotated_papers, note_extension="txt", edit=False):
         """Write annotations into pubs notes.
@@ -266,20 +344,18 @@ class ExtractPlugin(PapersPlugin):
         in the pubs notes directory. Creates new notes for
         citekeys missing a note or appends to existing.
         """
-        for contents in annotated_papers:
-            paper = contents[0]
-            annotations = contents[1]
+        for citekey, annotations in annotated_papers.items():
             if annotations:
-                notepath = self.broker.real_notepath(paper.citekey, note_extension)
+                notepath = self.broker.real_notepath(citekey, note_extension)
                 if check_file(notepath, fail=False):
                     self._append_to_note(notepath, annotations)
                 else:
                     self._write_new_note(notepath, annotations)
-                self.ui.info(f"Wrote annotations to {paper.citekey} note {notepath}.")
+                self.ui.info(f"Wrote annotations to {citekey} note {notepath}.")
 
                 if edit is True:
                     self.ui.edit_file(notepath, temporary=False)
-                NoteEvent(paper.citekey).send()
+                NoteEvent(citekey).send()
 
     def _write_new_note(self, notepath, annotations):
         """Create a new note containing the annotations.
@@ -289,7 +365,7 @@ class ExtractPlugin(PapersPlugin):
         """
         output = "# Annotations\n\n"
         for annotation in annotations:
-            output += f"{annotation}\n\n"
+            output += f"{annotation.format(self.formatting)}\n\n"
         write_file(notepath, output, "w")
 
     def _append_to_note(self, notepath, annotations):
@@ -300,13 +376,15 @@ class ExtractPlugin(PapersPlugin):
         """
         existing = read_text_file(notepath)
         # removed annotations already found in the note
-        existing_dropped = [x for x in annotations if x not in existing]
+        existing_dropped = [
+            x for x in annotations if x.format(self.formatting) not in existing
+        ]
         if not existing_dropped:
             return
 
         output = ""
         for annotation in existing_dropped:
-            output += f"{annotation}\n\n"
+            output += f"{annotation.format(self.formatting)}\n\n"
         write_file(notepath, output, "a")