diff --git a/src/streamer/parse/attach_markdown.py b/src/streamer/parse/attach_markdown.py index 85b5d6f..de2ae24 100644 --- a/src/streamer/parse/attach_markdown.py +++ b/src/streamer/parse/attach_markdown.py @@ -1,22 +1,30 @@ from __future__ import annotations +from pydantic import BaseModel + from streamer.parse.shard import Shard, StreamFile -class ShardWithMarkdown(Shard): - children: list[ShardWithMarkdown] +class ShardWithMarkdown(BaseModel): + markers: list[str] = [] + tags: list[str] = [] + start_line: int + end_line: int + children: list["ShardWithMarkdown"] = [] markdown_content: str -class StreamFileWithMarkdown(StreamFile): +class StreamFileWithMarkdown(BaseModel): + file_name: str shard: ShardWithMarkdown | None = None def attach_markdown_shard(shard: Shard, markdown_text: str) -> ShardWithMarkdown: lines = markdown_text.splitlines() markdown_content = "\n".join(lines[shard.start_line - 1 : shard.end_line]) + return ShardWithMarkdown( - **shard.model_dump(exclude=["children"]), + **shard.model_dump(exclude={"children"}), children=[ attach_markdown_shard(child, markdown_text) for child in shard.children ],