From 8879b909c00fa5299793ce90db8c1ee81d2af085 Mon Sep 17 00:00:00 2001
From: Valentin Berlier <berlier.v@gmail.com>
Date: Sun, 21 Feb 2021 08:58:17 +0100
Subject: [PATCH] fix: change document methods and add tests

---
 lectern/document.py    | 49 ++++++++++++++++++++--------
 tests/test_document.py | 72 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+), 14 deletions(-)
 create mode 100644 tests/test_document.py

diff --git a/lectern/document.py b/lectern/document.py
index ec409b7..b52fde9 100644
--- a/lectern/document.py
+++ b/lectern/document.py
@@ -5,9 +5,9 @@
 
 from dataclasses import InitVar, dataclass, field
 from pathlib import Path
-from typing import MutableMapping, Optional
+from typing import Any, Dict, Literal, MutableMapping, Optional, Tuple, overload
 
-from beet import Context, DataPack, ResourcePack
+from beet import Context, DataPack, File, ResourcePack
 from beet.core.utils import FileSystemPath, extra_field
 
 from .directive import Directive, get_builtin_directives
@@ -20,6 +20,9 @@ class Document:
     """Class representing a lectern document."""
 
     ctx: InitVar[Optional[Context]] = None
+    text: InitVar[Optional[str]] = None
+    markdown: InitVar[Optional[str]] = None
+    files: InitVar[Optional[FileSystemPath]] = None
 
     assets: ResourcePack = field(default_factory=ResourcePack)
     data: DataPack = field(default_factory=DataPack)
@@ -38,34 +41,53 @@ class Document:
         default_factory=MarkdownSerializer
     )
 
-    def __post_init__(self, ctx: Optional[Context]):
+    def __post_init__(
+        self,
+        ctx: Optional[Context],
+        text: Optional[str] = None,
+        markdown: Optional[str] = None,
+        files: Optional[FileSystemPath] = None,
+    ):
         if ctx:
             self.assets = ctx.assets
             self.data = ctx.data
+        if text:
+            self.add_text(text)
+        if markdown:
+            self.add_markdown(markdown, files)
 
     def add_text(self, source: str):
         """Extract pack fragments from plain text."""
         assets, data = self.text_extractor.extract(source, self.directives)
-
         self.assets.merge(assets)
         self.data.merge(data)
 
     def add_markdown(self, source: str, files: Optional[FileSystemPath] = None):
         """Extract pack fragments from markdown."""
         assets, data = self.markdown_extractor.extract(source, self.directives, files)
-
         self.assets.merge(assets)
         self.data.merge(data)
 
-    def serialize(self, as_text: bool = False) -> str:
-        """Serialize the inner data pack and the resource pack."""
-        return (
-            self.text_serializer.serialize(self.assets, self.data)
-            if as_text
-            else self.markdown_serializer.serialize_and_emit_files(
-                self.assets, self.data
-            )[0]
+    def get_text(self) -> str:
+        """Turn the data pack and the resource pack into text."""
+        return self.text_serializer.serialize(self.assets, self.data)
+
+    @overload
+    def get_markdown(
+        self, emit_files: Literal[True]
+    ) -> Tuple[str, Dict[str, File[Any, Any]]]:
+        ...
+
+    @overload
+    def get_markdown(self, emit_files: Literal[False] = False) -> str:
+        ...
+
+    def get_markdown(self, emit_files: bool = False):
+        """Turn the data pack and the resource pack into markdown."""
+        data, files = self.markdown_serializer.serialize_and_emit_files(
+            self.assets, self.data
         )
+        return (data, files) if emit_files else data
 
     def save(self, path: FileSystemPath, files: Optional[FileSystemPath] = None):
         """Save the serialized document at the specified location."""
@@ -84,5 +106,4 @@ def save(self, path: FileSystemPath, files: Optional[FileSystemPath] = None):
             data = self.text_serializer.serialize(self.assets, self.data)
 
         path.parent.mkdir(parents=True, exist_ok=True)
-
         path.write_text(data)
diff --git a/tests/test_document.py b/tests/test_document.py
new file mode 100644
index 0000000..16ced99
--- /dev/null
+++ b/tests/test_document.py
@@ -0,0 +1,72 @@
+import pytest
+from beet import DataPack, Function, ResourcePack
+
+from lectern import Document, InvalidFragment
+
+
+def test_empty():
+    assert Document() == Document()
+    assert Document().data == DataPack()
+    assert Document().assets == ResourcePack()
+    assert Document().get_text() == ""
+    assert Document().get_markdown() == ""
+    assert Document().get_markdown(emit_files=True) == ("", {})
+
+
+def test_data_pack():
+    pack = DataPack()
+    doc = Document(data=pack)
+    assert doc.data is pack
+
+
+def test_resource_pack():
+    pack = ResourcePack()
+    doc = Document(assets=pack)
+    assert doc.assets is pack
+
+
+def test_text_basic():
+    source = "@function demo:foo\nsay hello\n"
+    assert source in Document(text=source).get_text()
+
+
+def test_text_function():
+    pack = DataPack()
+    pack["demo:foo"] = Function(["say foo"])
+
+    doc = Document(data=pack)
+    doc.add_text("@function demo:bar\nsay bar\n")
+
+    assert pack.functions == {
+        "demo:foo": Function(["say foo"]),
+        "demo:bar": Function(["say bar"]),
+    }
+
+
+def test_text_tricky():
+    doc = Document(
+        text="some preamble\n\n"
+        "@function demo:foo\n"
+        "say foo\n"
+        "@other_thing hello world\n"
+        "say after\n"
+        " @function not taken into account\n"
+        "say next\n"
+        "@function demo:bar\n"
+        "say bar\n"
+    )
+    assert len(doc.data.functions) == 2
+
+
+def test_missing_argument():
+    with pytest.raises(
+        InvalidFragment, match="Missing directive argument 'full_name'."
+    ):
+        Document(text="@function\nsay hello")
+
+
+def test_extra_argument():
+    with pytest.raises(
+        InvalidFragment, match="Unexpected directive argument 'banana'."
+    ):
+        Document(text="@function demo:foo banana\nsay hello")