From 04c8a8ed0baf2eddbc63dff9a0c44a76b5fb350e Mon Sep 17 00:00:00 2001 From: Marty Oehme Date: Sun, 25 Dec 2022 12:17:05 +0100 Subject: [PATCH] Refactor extraction to use PaperAnnotated class --- README.md | 1 + extract/annotation.py | 30 +++++++++++++++++++------ extract/extract.py | 51 ++++++++++++++++++++++++------------------- 3 files changed, 52 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 03322d6..ab0f100 100644 --- a/README.md +++ b/README.md @@ -160,6 +160,7 @@ content, because then we can just use that. It is harder to parse if it does not - [ ] 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 + - [ ] use similarity search instead of literal search for existing annotation (levenshtein) - [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 diff --git a/extract/annotation.py b/extract/annotation.py index 3df7743..0baefbb 100644 --- a/extract/annotation.py +++ b/extract/annotation.py @@ -19,11 +19,33 @@ COLORS = { } +class PaperAnnotated(Paper): + def __init__(self, citekey, bibdata, metadata=None, annotations=[]): + super(PaperAnnotated, self).__init__(citekey, bibdata, metadata) + self.annotations = annotations + + @classmethod + def from_paper(cls, paper, annotations=[]): + return cls(paper.citekey, paper.bibdata, paper.metadata, annotations) + + def __repr__(self): + return "PaperAnnotated(%s, %s, %s)" % ( + self.citekey, + self.bibdata, + self.metadata, + ) + + def headline(self, short=False, max_authors=3): + headline = pretty.paper_oneliner( + self, citekey_only=short, max_authors=max_authors + ) + return re.sub(r"\[pdf\]", "", headline).rstrip() + + @dataclass class Annotation: """A PDF annotation object""" - paper: Paper file: str type: str = "Highlight" text: str = "" @@ -79,12 +101,6 @@ class Annotation: nearest = name return nearest - def headline(self, short=False, max_authors=3): - headline = pretty.paper_oneliner( - self.paper, citekey_only=short, max_authors=max_authors - ) - return re.sub(r"\[pdf\]", "", headline).rstrip() - def _color_similarity_ratio(self, color_one, color_two): """Return the similarity of two colors between 0 and 1. diff --git a/extract/extract.py b/extract/extract.py index cdcfdf9..d0496f4 100644 --- a/extract/extract.py +++ b/extract/extract.py @@ -10,7 +10,12 @@ from pubs import repo, pretty 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 -from .annotation import Annotation, COLOR_SIMILARITY_MINIMUM, TEXT_SIMILARITY_MINIMUM +from .annotation import ( + PaperAnnotated, + Annotation, + COLOR_SIMILARITY_MINIMUM, + TEXT_SIMILARITY_MINIMUM, +) CONFIRMATION_PAPER_THRESHOLD = 5 @@ -34,6 +39,7 @@ class ExtractPlugin(PapersPlugin): def __init__(self, conf, ui): self.ui = ui self.note_extension = conf["main"]["note_extension"] + self.max_authors = conf["main"]["max_authors"] self.repository = repo.Repository(conf) self.pubsdir = os.path.expanduser(conf["main"]["pubsdir"]) self.broker = self.repository.databroker @@ -123,12 +129,12 @@ 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: - annotations = self._get_annotations(file, paper) - papers_annotated[paper.citekey] = annotations + annotations = self._get_annotations(file) + papers_annotated.append(PaperAnnotated.from_paper(paper, annotations)) except fitz.FileDataError as e: self.ui.error(f"Document {file} is broken: {e}") return papers_annotated @@ -166,7 +172,7 @@ class ExtractPlugin(PapersPlugin): self.ui.message( "\n".join( pretty.paper_oneliner( - p, citekey_only=False, max_authors=conf["main"]["max_authors"] + p, citekey_only=False, max_authors=self.max_authors ) for p in papers ) @@ -188,7 +194,7 @@ class ExtractPlugin(PapersPlugin): self.ui.warning(f"{paper.citekey} has no valid document.") return path - def _get_annotations(self, filename, paper): + def _get_annotations(self, filename): """Extract annotations from a file. Returns all readable annotations contained in the file @@ -202,7 +208,6 @@ class ExtractPlugin(PapersPlugin): quote, note = self._retrieve_annotation_content(page, annot) a = Annotation( file=filename, - paper=paper, text=quote, content=note, colors=annot.colors, @@ -237,18 +242,18 @@ class ExtractPlugin(PapersPlugin): # highlight with selection not in note return (written, "") - def _to_stdout(self, annotated_papers, short_header=True): + def _to_stdout(self, annotated_papers, short_header=False): """Write annotations to stdout. Simply outputs the gathered annotations over stdout ready to be passed on through pipelines etc. """ output = "" - for citekey, annotations in annotated_papers.items(): + for paper in annotated_papers: output += ( - f"\n------ {annotations[0].headline(short=short_header)} ------\n\n" + f"\n------ {paper.headline(self.short_header, self.max_authors)} ------\n\n" ) - for annotation in annotations: + for annotation in paper.annotations: output += f"{annotation.format(self.formatting)}\n" output += "\n" self.ui.message(output.strip()) @@ -260,31 +265,31 @@ class ExtractPlugin(PapersPlugin): in the pubs notes directory. Creates new notes for citekeys missing a note or appends to existing. """ - for citekey, annotations in annotated_papers.items(): - if annotations: - notepath = self.broker.real_notepath(citekey, note_extension) + for paper in annotated_papers: + if paper.annotations: + notepath = self.broker.real_notepath(paper.citekey, note_extension) if check_file(notepath, fail=False): - self._append_to_note(notepath, annotations) + self._append_to_note(notepath, paper) else: - self._write_new_note(notepath, annotations, self.short_header) - self.ui.info(f"Wrote annotations to {citekey} note {notepath}.") + self._write_new_note(notepath, paper, paper.headline(short=True, max_authors=self.max_authors)) + self.ui.info(f"Wrote annotations to {paper.citekey} note {notepath}.") if edit is True: self.ui.edit_file(notepath, temporary=False) - NoteEvent(citekey).send() + NoteEvent(paper.citekey).send() - def _write_new_note(self, notepath, annotations, short_header): + def _write_new_note(self, notepath, paper, headline): """Create a new note containing the annotations. Will create a new note in the notes folder of pubs and fill it with the annotations extracted from pdf. """ - output = f"# {annotations[0].headline(short=short_header)}\n\n" - for annotation in annotations: + output = f"# {headline}\n\n" + for annotation in paper.annotations: output += f"{annotation.format(self.formatting)}\n\n" write_file(notepath, output, "w") - def _append_to_note(self, notepath, annotations): + def _append_to_note(self, notepath, paper): """Append new annotations to the end of a note. Looks through note to determine any new annotations which should be @@ -293,7 +298,7 @@ class ExtractPlugin(PapersPlugin): existing = read_text_file(notepath) # removed annotations already found in the note existing_dropped = [ - x for x in annotations if x.format(self.formatting) not in existing + x for x in paper.annotations if x.format(self.formatting) not in existing ] if not existing_dropped: return