84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
from typing import Generic, Iterable
|
|
|
|
from bs4 import BeautifulSoup, Tag
|
|
|
|
from .core import C, ContextParser, ParseError, Parser, PrepareParser, R
|
|
from .util import set_dict_nested_value
|
|
|
|
|
|
def _resolve_tag(source: Tag, selector: str | None = None) -> Tag:
|
|
if selector:
|
|
tag = source.select_one(selector)
|
|
if tag is None:
|
|
raise ParseError("Tag not found", selector)
|
|
return tag
|
|
return source
|
|
|
|
|
|
class _HTMLParser(Parser[str, BeautifulSoup]):
|
|
def parse(self, data: str) -> BeautifulSoup:
|
|
return BeautifulSoup(data, features="html.parser")
|
|
|
|
|
|
HTML_PARSER = _HTMLParser()
|
|
|
|
|
|
class HTMLParser(PrepareParser[str, Tag, R], Generic[R]):
|
|
selector: str | None = None
|
|
|
|
def prepare(self, data: str) -> Tag:
|
|
return _resolve_tag(HTML_PARSER.parse(data), self.selector)
|
|
|
|
|
|
class HTMLContextParser(ContextParser[C, Tag, R], Generic[C, R]):
|
|
selector: str | None = None
|
|
|
|
def prepare_context(self, data: str) -> Tag:
|
|
return _resolve_tag(HTML_PARSER.parse(data), self.selector)
|
|
|
|
|
|
class TagParser(PrepareParser[Tag, Tag, R], Generic[R]):
|
|
selector: str | None = None
|
|
|
|
def prepare(self, data: Tag) -> Tag:
|
|
return _resolve_tag(data, self.selector)
|
|
|
|
|
|
class TagIterParser(Parser[Tag, Iterable[R]], Generic[R]):
|
|
selector: str
|
|
|
|
def extract(self, data: Tag) -> R:
|
|
raise NotImplementedError
|
|
|
|
def parse(self, data: Tag) -> Iterable[R]:
|
|
tags = data.select(self.selector)
|
|
for tag in tags:
|
|
yield self.extract(tag)
|
|
|
|
|
|
class TableValueParser(TagIterParser[R], Generic[R]):
|
|
model_key: str | list[str]
|
|
|
|
def extract(self, data: Tag) -> R:
|
|
raise NotImplementedError
|
|
|
|
|
|
class TableParser(TagParser[list[R]], Generic[R]):
|
|
value_parsers: list[TableValueParser]
|
|
|
|
def create_model(self, data: dict) -> R:
|
|
raise NotImplementedError
|
|
|
|
def extract(self, data: Tag) -> list[R]:
|
|
results = []
|
|
for parser in self.value_parsers:
|
|
for index, value in enumerate(parser.parse(data)):
|
|
while index > len(results) - 1:
|
|
results.append({})
|
|
if isinstance(parser.model_key, list):
|
|
for model_key_index, model_key in enumerate(parser.model_key):
|
|
set_dict_nested_value(results[index], model_key, value[model_key_index])
|
|
else:
|
|
set_dict_nested_value(results[index], parser.model_key, value)
|
|
return [self.create_model(item) for item in results]
|