Update src/text.py

src/text.py

@@ -0,0 +1,130 @@
+from markdownify import markdownify as md
+from bs4 import BeautifulSoup as BS
+from IPython.display import display, Markdown
+from urllib.parse import urljoin
+from newspaper import Article
+import re
+import markdown
+
+
+def clean(s):
+    s = s.replace("\t", "\\t")
+    s = s.replace("\n", "\\n")
+    return s
+
+class DocTree:
+    def __init__(self, content):
+        self.content = content
+        self.max_depth = 6
+
+    def get_sections(self, *location_ids):
+        out = self.content
+        for id_ in location_ids:
+            out = out[id_]
+        return out
+    
+    def merge_sections(self, elems):
+        if not isinstance(elems[0], list):
+            return '\n\n '.join(elems)
+        out = []
+        for e in elems:
+            out.append(self.merge_sections(e))
+        return '\n\n '.join(map(clean, out))
+
+    def get_merged_sections(self, *location_ids):
+        return [self.merge_sections(s) for s in self.get_sections(*location_ids)]
+
+    def as_markdown(self, content):
+        return md(content)
+
+    def get_sections_by_depth(self, depth):
+        return self._get_sections_by_depth(self.content, depth)
+
+    @staticmethod
+    def _get_sections_by_depth(content, depth):
+        """Returns a list of merged sections at a specific depth"""
+        if depth == 0:
+            return content
+        out = []
+        for elem in content:
+            out += DocTree._get_sections_by_depth(elem, depth - 1)
+        return out
+
+
+def fix_relative_links(url, article_content):
+    if 'http' in url:
+        base_url = '/'.join(url.split('/')[:3])
+    else:
+        base_url = url.split('/')
+    pat = re.compile(r'\[(.*?)\]\((.*?)\)', flags=re.IGNORECASE)
+    res = pat.findall(article_content)
+    if res:
+        for g in res:
+            url = urljoin(base_url, g[1]) if g[1].startswith('/') else g[1]
+            article_content = article_content.replace(f'[{g[0]}]({g[1]})', f'[{g[0]}]({url})')
+    else:print('not found')
+    return article_content
+
+
+def extract_article(url):
+    article = Article(url)
+    article.download()
+    article.parse()
+    return article
+
+
+def select_content(html_code, elem_class, class_name):
+    print(f'Calling select_content with {elem_class}, {class_name}')
+    if class_name.startswith('.'):
+        class_name = class_name[1:]
+        elem_id = None
+    elif class_name.startswith('#'):
+        elem_id = class_name[1:]
+        class_name = None
+    else:
+        elem_id = None
+        class_name = None
+    return md(str(BS(html_code, features="lxml").find(elem_class, class_=class_name, id=elem_id)))
+
+
+def split_by_heading(html_content, _i):
+    if _i >= 7:
+        return html_content
+    elems = []
+    for idx, elem in enumerate([i for i in html_content.split(f'<h{_i}') if i]):
+        if idx > 0 or elem.startswith('>'):
+            elem = f'<h{_i}{elem}'
+        elems.append(split_by_heading(elem, _i+1))
+    return elems
+
+def doctree_from_url(url, elem_class='div', class_name='article-body'):
+    article = extract_article(url)
+    # convert to MD to handle splitting better
+    article_content = select_content(article.html, elem_class, class_name)
+    article_content = (f"# {article.title}\n\n" + article_content).replace('\n\n', '\n').replace('#', '%%@@%%')
+    # fix relative website links
+    article_content = fix_relative_links(url, article_content)
+    # convert back to HTML
+    html_content = markdown.markdown(article_content).replace('%%@@%%', '#')
+    doc_tree = DocTree(split_by_heading(html_content, 1))
+
+    #assert doc_tree.merge_sections(doc_tree.get_sections(0)).replace('\n', '').replace(html_content.replace('\n', ''), '') == '', 'Document inconsistent. Manual adjustments required.'
+    return doc_tree
+
+
+def get_selectors_for_class(url, elem_class):
+    article = extract_article(url)
+
+    html_content = article.html
+    soup = BS(html_content, features="lxml")
+    classes = set()
+    ids = set()
+    for elem in soup.find_all(elem_class):
+        if elem.get('class'):
+            for c in elem.get('class'):
+                classes |= {f".{c}"}
+        if elem.get('id'):
+            for c in elem.get('id'):
+                ids |= {f"#{c}"}
+
+    return ids | classes