From 090a49064f276f30ede53e8a87c9b7fd0d4c70bf Mon Sep 17 00:00:00 2001
From: Oleh Prypin <oleh@pryp.in>
Date: Sun, 21 Aug 2022 16:32:13 +0200
Subject: [PATCH] Type fixes according to mypy

---
 mkdocs_section_index/__init__.py | 10 ++++++----
 mkdocs_section_index/plugin.py   |  2 ++
 mkdocs_section_index/rewrites.py |  9 +++++++--
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/mkdocs_section_index/__init__.py b/mkdocs_section_index/__init__.py
index 213260a..541359a 100644
--- a/mkdocs_section_index/__init__.py
+++ b/mkdocs_section_index/__init__.py
@@ -1,15 +1,17 @@
-from typing import Sequence
+from typing import List, Union
 
 from mkdocs.config import Config
 from mkdocs.structure.files import File
-from mkdocs.structure.nav import Section
+from mkdocs.structure.nav import Link, Section
 from mkdocs.structure.pages import Page
 
 __all__ = ["SectionPage"]
 
 
-class SectionPage(Section, Page):
-    def __init__(self, title: str, file: File, config: Config, children: Sequence):
+class SectionPage(Section, Page):  # type: ignore[misc]
+    def __init__(
+        self, title: str, file: File, config: Config, children: List[Union[Page, Section, Link]]
+    ):
         Page.__init__(self, title=title, file=file, config=config)
         Section.__init__(self, title=title, children=children)
         self.is_section = self.is_page = True
diff --git a/mkdocs_section_index/plugin.py b/mkdocs_section_index/plugin.py
index 956aab0..b775c96 100644
--- a/mkdocs_section_index/plugin.py
+++ b/mkdocs_section_index/plugin.py
@@ -31,6 +31,7 @@ def on_nav(self, nav: Navigation, config, files) -> Navigation:
                 if not page.title and page.url:
                     # The page becomes a section-page.
                     page.__class__ = SectionPage
+                    assert isinstance(page, SectionPage)
                     page.is_section = page.is_page = True
                     page.title = section.title
                     # The page leaves the section but takes over children that used to be its peers.
@@ -44,6 +45,7 @@ def on_nav(self, nav: Navigation, config, files) -> Navigation:
         return nav
 
     def on_env(self, env: Environment, config, files) -> Environment:
+        assert env.loader is not None
         env.loader = self._loader = rewrites.TemplateRewritingLoader(env.loader)
         return env
 
diff --git a/mkdocs_section_index/rewrites.py b/mkdocs_section_index/rewrites.py
index 848f7f2..9b4aaa9 100644
--- a/mkdocs_section_index/rewrites.py
+++ b/mkdocs_section_index/rewrites.py
@@ -1,7 +1,7 @@
 import logging
 import pathlib
 import textwrap
-from typing import Optional, Tuple
+from typing import Callable, Optional, Tuple
 
 import mkdocs.utils
 from jinja2 import BaseLoader, Environment
@@ -18,9 +18,13 @@ def __init__(self, loader: BaseLoader):
         self.loader = loader
         self.found_supported_theme = False
 
-    def get_source(self, environment: Environment, template: str) -> Tuple[str, str, bool]:
+    def get_source(
+        self, environment: Environment, template: str
+    ) -> Tuple[str, str, Optional[Callable[[], bool]]]:
+        src: Optional[str]
         src, filename, uptodate = self.loader.get_source(environment, template)
         old_src = src
+        assert filename is not None
         path = pathlib.Path(filename).as_posix()
 
         if path.endswith("/mkdocs/templates/sitemap.xml"):
@@ -51,6 +55,7 @@ def _transform_mkdocs_sitemap_template(src: str) -> Optional[str]:
             "{%- else %}",
             "{%- endif %}{% if item.url %}",
         )
+    return None
 
 
 def _transform_material_nav_item_template(src: str) -> str: