Trying to refactor the parsing

buster/docparser.py

@@ -1,5 +1,6 @@
 import glob
 import os
+from typing import Type
 
 import numpy as np
 import pandas as pd
@@ -42,7 +43,7 @@ supported_docs = {
 
 
 def get_all_documents(
-    root_dir: str, base_url: str, parser: Parser, min_section_length: int = 100, max_section_length: int = 2000
+    root_dir: str, base_url: str, parser_cls: Type[Parser], min_section_length: int = 100, max_section_length: int = 2000
 ) -> pd.DataFrame:
     """Parse all HTML files in `root_dir`, and extract all sections.
 
@@ -60,12 +61,12 @@ def get_all_documents(
             source = f.read()
 
         soup = BeautifulSoup(source, "html.parser")
-        soup_parser = parser(soup, base_url, file, min_section_length, max_section_length)
-        sections_file, urls_file, names_file = soup_parser.parse()
-
-        sections.extend(sections_file)
-        urls.extend(urls_file)
-        names.extend(names_file)
+        parser = parser_cls(soup, base_url, file, min_section_length, max_section_length)
+        # sections_file, urls_file, names_file =
+        for section in parser.parse():
+            sections.append(section.text)
+            urls.append(section.url)
+            names.append(section.name)
 
     documents_df = pd.DataFrame.from_dict({"title": names, "url": urls, "content": sections})
 

buster/parser.py

@@ -1,151 +1,127 @@
+from abc import ABC, abstractmethod
+from dataclasses import InitVar, dataclass, field
+from itertools import takewhile, zip_longest
 import math
 import os
+from typing import Generator
 
 import bs4
 import pandas as pd
 from bs4 import BeautifulSoup
 
 
-def parse_section(nodes: list[bs4.element.NavigableString]) -> str:
-    section = []
-    for node in nodes:
-        if node.name == "table":
-            node_text = pd.read_html(node.prettify())[0].to_markdown(index=False, tablefmt="github")
-        elif node.name == "script":
-            continue
-        else:
-            node_text = node.text
-        section.append(node_text)
-    section = "".join(section)
-
-    return section
-
-
-class Parser:
-    def __init__(
-        self,
-        soup: BeautifulSoup,
-        base_url: str,
-        filename: str,
-        min_section_length: int = 100,
-        max_section_length: int = 2000,
-    ):
-        self.soup = soup
-        self.base_url = base_url
-        self.filename = filename
-        self.min_section_length = min_section_length
-        self.max_section_length = max_section_length
-
-    def parse(self) -> tuple[list[str], list[str], list[str]]:
+@dataclass
+class Section:
+    url: str
+    name: str
+    nodes: InitVar[list[bs4.element.NavigableString]]
+    text: str = field(init=False)
+
+    def __post_init__(self, nodes: list[bs4.element.NavigableString]):
+        section = []
+        for node in nodes:
+            if node.name == "table":
+                node_text = pd.read_html(node.prettify())[0].to_markdown(index=False, tablefmt="github")
+            elif node.name == "script":
+                continue
+            else:
+                node_text = node.text
+            section.append(node_text)
+        self.text = "".join(section).strip()
+
+    def __len__(self) -> int:
+        return len(self.text)
+
+    @classmethod
+    def from_text(cls, text: str, url: str, name: str) -> 'Section':
+        """Alternate constructor, without parsing."""
+        section = cls.__new__(cls)  # Allocate memory, does not call __init__
+        # Does the init here.
+        section.text = text
+        section.url = url
+        section.name = name
+
+        return section
+
+    def get_chunks(self, min_length: int, max_length: int) -> Generator['Section', None, None]:
+        """Split a section into chunks."""
+        if len(self) > max_length:
+            # Get the number of chunk, by dividing and rounding up.
+            # Then, split the section into equal lenght chunks.
+            # This could results in chunks below the minimum length,
+            # and will truncate the end of the section.
+            n_chunks = (len(self) + max_length - 1) // max_length
+            length = len(self) // n_chunks
+            for chunk in range(n_chunks):
+                start = chunk * length
+                yield Section.from_text(
+                    self.text[start: start + length],
+                    self.url, self.name)
+        elif len(self) > min_length:
+            yield self
+        return
+
+
+@dataclass
+class Parser(ABC):
+    soup: BeautifulSoup
+    base_url: str
+    filename: str
+    min_section_length: int = 100
+    max_section_length: int = 2000
+
+    @abstractmethod
+    def build_url(self, suffix: str) -> str:
         ...
 
-    def find_sections(self) -> bs4.element.ResultSet:
+    @abstractmethod
+    def find_sections(self) -> Generator[Section, None, None]:
         ...
 
-    def build_url(self, suffix: str) -> str:
-        ...
+    def parse(self) -> list[Section]:
+        """Parse the documents into sections, respecting the lenght constraints."""
+        sections = []
+        for section in self.find_sections():
+            sections.extend(section.get_chunks(self.min_section_length, self.max_section_length))
+        return sections
 
 
 class SphinxParser(Parser):
-    def parse(self) -> tuple[list[str], list[str], list[str]]:
-        found = self.find_sections()
 
-        sections = []
-        urls = []
-        names = []
-        for i in range(len(found)):
-            section_found = found[i]
+    def find_sections(self) -> Generator[Section, None, None]:
+        for section in self.soup.find_all("a", href=True, class_="headerlink"):
+            container = section.parent.parent
+            section_href = container.find_all("a", href=True, class_="headerlink")
 
-            section_soup = section_found.parent.parent
-            section_href = section_soup.find_all("a", href=True, class_="headerlink")
+            url = self.build_url(section["href"].strip().replace("\n", ""))
+            name = section.parent.text
 
             # If sections has subsections, keep only the part before the first subsection
-            if len(section_href) > 1 and section_soup.section is not None:
-                section_siblings = list(section_soup.section.previous_siblings)[::-1]
-                section = parse_section(section_siblings)
+            if len(section_href) > 1 and container.section is not None:
+                siblings = list(container.section.previous_siblings)[::-1]
+                section = Section(url, name, siblings)
             else:
-                section = parse_section(section_soup.children)
-
-            # Remove special characters, plus newlines in some url and section names.
-            section = section.strip()
-            url = section_found["href"].strip().replace("\n", "")
-            name = section_found.parent.text.strip()[:-1].replace("\n", "")
-
-            url = self.build_url(url)
-
-            # If text is too long, split into chunks of equal sizes
-            if len(section) > self.max_section_length:
-                n_chunks = math.ceil(len(section) / float(self.max_section_length))
-                separator_index = math.floor(len(section) / n_chunks)
-
-                section_chunks = [section[separator_index * i : separator_index * (i + 1)] for i in range(n_chunks)]
-                url_chunks = [url] * n_chunks
-                name_chunks = [name] * n_chunks
-
-                sections.extend(section_chunks)
-                urls.extend(url_chunks)
-                names.extend(name_chunks)
-            # If text is not too short, add in 1 chunk
-            elif len(section) > self.min_section_length:
-                sections.append(section)
-                urls.append(url)
-                names.append(name)
-
-        return sections, urls, names
-
-    def find_sections(self) -> bs4.element.ResultSet:
-        return self.soup.find_all("a", href=True, class_="headerlink")
+                section = Section(url, name, container.children)
+            yield section
+        return
 
     def build_url(self, suffix: str) -> str:
         return self.base_url + self.filename + suffix
 
 
 class HuggingfaceParser(Parser):
-    def parse(self) -> tuple[list[str], list[str], list[str]]:
-        found = self.find_sections()
 
-        sections = []
-        urls = []
-        names = []
-        for i in range(len(found)):
-            section_href = found[i].find("a", href=True, class_="header-link")
-
-            section_nodes = []
-            for element in found[i].find_next_siblings():
-                if i + 1 < len(found) and element == found[i + 1]:
-                    break
-                section_nodes.append(element)
-            section = parse_section(section_nodes)
-
-            # Remove special characters, plus newlines in some url and section names.
-            section = section.strip()
-            url = section_href["href"].strip().replace("\n", "")
-            name = found[i].text.strip().replace("\n", "")
-
-            url = self.build_url(url)
-
-            # If text is too long, split into chunks of equal sizes
-            if len(section) > self.max_section_length:
-                n_chunks = math.ceil(len(section) / float(self.max_section_length))
-                separator_index = math.floor(len(section) / n_chunks)
-
-                section_chunks = [section[separator_index * i : separator_index * (i + 1)] for i in range(n_chunks)]
-                url_chunks = [url] * n_chunks
-                name_chunks = [name] * n_chunks
-
-                sections.extend(section_chunks)
-                urls.extend(url_chunks)
-                names.extend(name_chunks)
-            # If text is not too short, add in 1 chunk
-            elif len(section) > self.min_section_length:
-                sections.append(section)
-                urls.append(url)
-                names.append(name)
-
-        return sections, urls, names
-
-    def find_sections(self) -> bs4.element.ResultSet:
-        return self.soup.find_all(["h1", "h2", "h3"], class_="relative group")
+    def find_sections(self) -> Generator[Section, None, None]:
+        sections = self.soup.find_all(["h1", "h2", "h3"], class_="relative group")
+        for section, next_section in zip_longest(sections, sections[1:]):
+            href = section.find("a", href=True, class_="header-link")
+            nodes = list(takewhile(lambda sibling: sibling != next_section, section.find_next_siblings()))
+
+            url = self.build_url(href["href"].strip().replace("\n", ""))
+            name = section.text.strip().replace("\n", "")
+            yield Section(url, name, nodes)
+
+        return
 
     def build_url(self, suffix: str) -> str:
         # The splitext is to remove the .html extension