diff --git a/extract/extract.py b/extract/extract.py index 18f3f0f..8c95c63 100644 --- a/extract/extract.py +++ b/extract/extract.py @@ -1,6 +1,7 @@ import os import re import argparse +import math from dataclasses import dataclass from typing import Tuple @@ -18,6 +19,15 @@ from pubs.query import get_paper_filter CONFIRMATION_PAPER_THRESHOLD = 5 +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: @@ -55,6 +65,17 @@ class Annotation: ) return pattern.sub(lambda x: replacements[x.group(0)], output) + @property + def colorname(self): + annot_colors = self.colors.get("stroke") or self.colors.get("fill") + nearest = None + smallest_dist = 2.0 + for name, values in COLORS.items(): + dist = math.dist([*values], [*annot_colors]) + if dist < smallest_dist: + smallest_dist = dist + nearest = name + return nearest class ExtractPlugin(PapersPlugin): """Extract annotations from any pdf document. @@ -86,6 +107,7 @@ class ExtractPlugin(PapersPlugin): "formatting", "{newline}{quote_begin}> {quote} {quote_end}[{page}]{note_begin}{newline}Note: {note}{note_end}", ) + self.color_mapping = settings.get("color_mapping", {}) def update_parser(self, subparsers, _): """Allow the usage of the pubs extract subcommand""" @@ -167,6 +189,9 @@ class ExtractPlugin(PapersPlugin): self.ui.error(f"Document {file} is broken: {e}") return papers_annotated + def mapped_tag(self, colorname): + return self.color_mapping.get(colorname) + def _gather_papers(self, conf, args): """Get all papers for citekeys.