Add types to codebase

.github/workflows/test.yml

@@ -27,3 +27,6 @@ jobs:
       - name: Run tests
         run: |
           pytest . --doctest-modules --doctest-glob "README.md"
+      - name: Run type checking
+        run: |
+          mypy .

frontmatter/__init__.py

@@ -2,15 +2,20 @@
 """
 Python Frontmatter: Parse and manage posts with YAML frontmatter
 """
+from __future__ import annotations
 
 import codecs
-import re
-
+import io
+from typing import TYPE_CHECKING, Iterable
 
 from .util import u
 from .default_handlers import YAMLHandler, JSONHandler, TOMLHandler
 
 
+if TYPE_CHECKING:
+    from .default_handlers import BaseHandler
+
+
 __all__ = ["parse", "load", "loads", "dump", "dumps"]
 
 
@@ -22,7 +27,7 @@
 ]
 
 
-def detect_format(text, handlers):
+def detect_format(text: str, handlers: Iterable[BaseHandler]) -> BaseHandler | None:
     """
     Figure out which handler to use, based on metadata.
     Returns a handler instance or None.
@@ -40,7 +45,12 @@ def detect_format(text, handlers):
     return None
 
 
-def parse(text, encoding="utf-8", handler=None, **defaults):
+def parse(
+    text: str,
+    encoding: str = "utf-8",
+    handler: BaseHandler | None = None,
+    **defaults: object,
+) -> tuple[dict[str, object], str]:
     """
     Parse text with frontmatter, return metadata and content.
     Pass in optional metadata defaults as keyword args.
@@ -79,14 +89,14 @@ def parse(text, encoding="utf-8", handler=None, **defaults):
         return metadata, text
 
     # parse, now that we have frontmatter
-    fm = handler.load(fm)
-    if isinstance(fm, dict):
-        metadata.update(fm)
+    fm_data = handler.load(fm)
+    if isinstance(fm_data, dict):
+        metadata.update(fm_data)
 
     return metadata, content.strip()
 
 
-def check(fd, encoding="utf-8"):
+def check(fd: str | io.IOBase, encoding: str = "utf-8") -> bool:
     """
     Check if a file-like object or filename has a frontmatter,
     return True if exists, False otherwise.
@@ -109,7 +119,7 @@ def check(fd, encoding="utf-8"):
     return checks(text, encoding)
 
 
-def checks(text, encoding="utf-8"):
+def checks(text: str, encoding: str = "utf-8") -> bool:
     """
     Check if a text (binary or unicode) has a frontmatter,
     return True if exists, False otherwise.
@@ -127,7 +137,12 @@ def checks(text, encoding="utf-8"):
     return detect_format(text, handlers) != None
 
 
-def load(fd, encoding="utf-8", handler=None, **defaults):
+def load(
+    fd: str | io.IOBase,
+    encoding: str = "utf-8",
+    handler: BaseHandler | None = None,
+    **defaults: object,
+) -> Post:
     """
     Load and parse a file-like object or filename,
     return a :py:class:`post <frontmatter.Post>`.
@@ -150,7 +165,12 @@ def load(fd, encoding="utf-8", handler=None, **defaults):
     return loads(text, encoding, handler, **defaults)
 
 
-def loads(text, encoding="utf-8", handler=None, **defaults):
+def loads(
+    text: str,
+    encoding: str = "utf-8",
+    handler: BaseHandler | None = None,
+    **defaults: object,
+) -> Post:
     """
     Parse text (binary or unicode) and return a :py:class:`post <frontmatter.Post>`.
 
@@ -166,7 +186,13 @@ def loads(text, encoding="utf-8", handler=None, **defaults):
     return Post(content, handler, **metadata)
 
 
-def dump(post, fd, encoding="utf-8", handler=None, **kwargs):
+def dump(
+    post: Post,
+    fd: str | io.IOBase,
+    encoding: str = "utf-8",
+    handler: BaseHandler | None = None,
+    **kwargs: object,
+) -> None:
     """
     Serialize :py:class:`post <frontmatter.Post>` to a string and write to a file-like object.
     Text will be encoded on the way out (utf-8 by default).
@@ -213,7 +239,7 @@ def dump(post, fd, encoding="utf-8", handler=None, **kwargs):
             f.write(content)
 
 
-def dumps(post, handler=None, **kwargs):
+def dumps(post: Post, handler: BaseHandler | None = None, **kwargs: object) -> str:
     """
     Serialize a :py:class:`post <frontmatter.Post>` to a string and return text.
     This always returns unicode text, which can then be encoded.
@@ -265,46 +291,48 @@ class Post(object):
     For convenience, metadata values are available as proxied item lookups.
     """
 
-    def __init__(self, content, handler=None, **metadata):
+    def __init__(
+        self, content: str, handler: BaseHandler | None = None, **metadata: object
+    ) -> None:
         self.content = str(content)
         self.metadata = metadata
         self.handler = handler
 
-    def __getitem__(self, name):
+    def __getitem__(self, name: str) -> object:
         "Get metadata key"
         return self.metadata[name]
 
-    def __contains__(self, item):
+    def __contains__(self, item: object) -> bool:
         "Check metadata contains key"
         return item in self.metadata
 
-    def __setitem__(self, name, value):
+    def __setitem__(self, name: str, value: object) -> None:
         "Set a metadata key"
         self.metadata[name] = value
 
-    def __delitem__(self, name):
+    def __delitem__(self, name: str) -> None:
         "Delete a metadata key"
         del self.metadata[name]
 
-    def __bytes__(self):
+    def __bytes__(self) -> bytes:
         return self.content.encode("utf-8")
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.content
 
-    def get(self, key, default=None):
+    def get(self, key: str, default: object = None) -> object:
         "Get a key, fallback to default"
         return self.metadata.get(key, default)
 
-    def keys(self):
+    def keys(self) -> Iterable[str]:
         "Return metadata keys"
         return self.metadata.keys()
 
-    def values(self):
+    def values(self) -> Iterable[object]:
         "Return metadata values"
         return self.metadata.values()
 
-    def to_dict(self):
+    def to_dict(self) -> dict[str, object]:
         "Post as a dict, for serializing"
         d = self.metadata.copy()
         d["content"] = self.content

frontmatter/conftest.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import pytest
 
 
 @pytest.fixture(autouse=True)
-def add_globals(doctest_namespace):
+def add_globals(doctest_namespace: dict[str, object]) -> None:
     import frontmatter
 
     doctest_namespace["frontmatter"] = frontmatter

frontmatter/default_handlers.py

@@ -116,11 +116,19 @@
 
 
 """
+from __future__ import annotations
 
 import json
 import re
 import yaml
 
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, Type
+
+SafeDumper: Type[yaml.CDumper] | Type[yaml.SafeDumper]
+SafeLoader: Type[yaml.CSafeLoader] | Type[yaml.SafeLoader]
+toml: ModuleType | None
+
 try:
     from yaml import CSafeDumper as SafeDumper
     from yaml import CSafeLoader as SafeLoader
@@ -136,6 +144,10 @@
 from .util import u
 
 
+if TYPE_CHECKING:
+    from frontmatter import Post
+
+
 __all__ = ["BaseHandler", "YAMLHandler", "JSONHandler"]
 
 if toml:
@@ -159,11 +171,16 @@ class BaseHandler:
     All default handlers are subclassed from BaseHandler.
     """
 
-    FM_BOUNDARY = None
-    START_DELIMITER = None
-    END_DELIMITER = None
+    FM_BOUNDARY: re.Pattern[str] | None = None
+    START_DELIMITER: str | None = None
+    END_DELIMITER: str | None = None
 
-    def __init__(self, fm_boundary=None, start_delimiter=None, end_delimiter=None):
+    def __init__(
+        self,
+        fm_boundary: re.Pattern[str] | None = None,
+        start_delimiter: str | None = None,
+        end_delimiter: str | None = None,
+    ):
         self.FM_BOUNDARY = fm_boundary or self.FM_BOUNDARY
         self.START_DELIMITER = start_delimiter or self.START_DELIMITER
         self.END_DELIMITER = end_delimiter or self.END_DELIMITER
@@ -176,38 +193,40 @@ def __init__(self, fm_boundary=None, start_delimiter=None, end_delimiter=None):
                 )
             )
 
-    def detect(self, text):
+    def detect(self, text: str) -> bool:
         """
         Decide whether this handler can parse the given ``text``,
         and return True or False.
 
         Note that this is *not* called when passing a handler instance to
         :py:func:`frontmatter.load <frontmatter.load>` or :py:func:`loads <frontmatter.loads>`.
         """
+        assert self.FM_BOUNDARY is not None
         if self.FM_BOUNDARY.match(text):
             return True
         return False
 
-    def split(self, text):
+    def split(self, text: str) -> tuple[str, str]:
         """
         Split text into frontmatter and content
         """
+        assert self.FM_BOUNDARY is not None
         _, fm, content = self.FM_BOUNDARY.split(text, 2)
         return fm, content
 
-    def load(self, fm):
+    def load(self, fm: str) -> dict[str, Any]:
         """
         Parse frontmatter and return a dict
         """
         raise NotImplementedError
 
-    def export(self, metadata, **kwargs):
+    def export(self, metadata: dict[str, object], **kwargs: object) -> str:
         """
         Turn metadata back into text
         """
         raise NotImplementedError
 
-    def format(self, post, **kwargs):
+    def format(self, post: Post, **kwargs: object) -> str:
         """
         Turn a post into a string, used in ``frontmatter.dumps``
         """
@@ -233,23 +252,23 @@ class YAMLHandler(BaseHandler):
     FM_BOUNDARY = re.compile(r"^-{3,}\s*$", re.MULTILINE)
     START_DELIMITER = END_DELIMITER = "---"
 
-    def load(self, fm, **kwargs):
+    def load(self, fm: str, **kwargs: object) -> Any:
         """
         Parse YAML front matter. This uses yaml.SafeLoader by default.
         """
         kwargs.setdefault("Loader", SafeLoader)
-        return yaml.load(fm, **kwargs)
+        return yaml.load(fm, **kwargs)  # type: ignore[arg-type]
 
-    def export(self, metadata, **kwargs):
+    def export(self, metadata: dict[str, object], **kwargs: object) -> str:
         """
         Export metadata as YAML. This uses yaml.SafeDumper by default.
         """
         kwargs.setdefault("Dumper", SafeDumper)
         kwargs.setdefault("default_flow_style", False)
         kwargs.setdefault("allow_unicode", True)
 
-        metadata = yaml.dump(metadata, **kwargs).strip()
-        return u(metadata)  # ensure unicode
+        metadata_str = yaml.dump(metadata, **kwargs).strip()  # type: ignore[call-overload]
+        return u(metadata_str)  # ensure unicode
 
 
 class JSONHandler(BaseHandler):
@@ -263,18 +282,18 @@ class JSONHandler(BaseHandler):
     START_DELIMITER = ""
     END_DELIMITER = ""
 
-    def split(self, text):
+    def split(self, text: str) -> tuple[str, str]:
         _, fm, content = self.FM_BOUNDARY.split(text, 2)
         return "{" + fm + "}", content
 
-    def load(self, fm, **kwargs):
-        return json.loads(fm, **kwargs)
+    def load(self, fm: str, **kwargs: object) -> Any:
+        return json.loads(fm, **kwargs)  # type: ignore[arg-type]
 
-    def export(self, metadata, **kwargs):
+    def export(self, metadata: dict[str, object], **kwargs: object) -> str:
         "Turn metadata into JSON"
         kwargs.setdefault("indent", 4)
-        metadata = json.dumps(metadata, **kwargs)
-        return u(metadata)
+        metadata_str = json.dumps(metadata, **kwargs)  # type: ignore[arg-type]
+        return u(metadata_str)
 
 
 if toml:
@@ -289,14 +308,15 @@ class TOMLHandler(BaseHandler):
         FM_BOUNDARY = re.compile(r"^\+{3,}\s*$", re.MULTILINE)
         START_DELIMITER = END_DELIMITER = "+++"
 
-        def load(self, fm, **kwargs):
+        def load(self, fm: str, **kwargs: object) -> Any:
+            assert toml is not None
             return toml.loads(fm, **kwargs)
 
-        def export(self, metadata, **kwargs):
+        def export(self, metadata: dict[str, object], **kwargs: object) -> str:
             "Turn metadata into TOML"
-            metadata = toml.dumps(metadata)
-            return u(metadata)
-
+            assert toml is not None
+            metadata_str = toml.dumps(metadata)
+            return u(metadata_str)
 
 else:
-    TOMLHandler = None
+    TOMLHandler: Type[TOMLHandler] | None = None  #  type: ignore[no-redef]

frontmatter/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561. This package uses inline types.