Split plugin flatten into two functions to fix type issues

Author: object-Object

src/hexdoc/cli/utils/info.py

@@ -12,7 +12,7 @@
     Properties,
 )
 from hexdoc.plugin import ModPlugin, PluginManager
-from hexdoc.plugin.manager import flatten, import_package
+from hexdoc.plugin.manager import flatten_hook_return, import_package
 
 logger = logging.getLogger(__name__)
 
@@ -53,7 +53,7 @@ def _plugins(pm: PluginManager):
 
 
 def _jinja_template_roots(mod_plugin: ModPlugin):
-    for package, folder in flatten([mod_plugin.jinja_template_root() or []]):
+    for package, folder in flatten_hook_return(mod_plugin.jinja_template_root()):
         module_path = _get_package_path(package)
         folder_path = module_path / folder
         yield _relative_path(folder_path)

src/hexdoc/plugin/manager.py

@@ -35,7 +35,7 @@
 from .book_plugin import BookPlugin
 from .mod_plugin import DefaultRenderedTemplates, ModPlugin, ModPluginWithBook
 from .specs import HEXDOC_PROJECT_NAME, PluginSpec
-from .types import HookReturns
+from .types import HookReturn, HookReturns
 
 _T = TypeVar("_T")
 
@@ -117,19 +117,19 @@ def init_plugins(self):
 
     def _init_book_plugins(self):
         caller = self._hook_caller(PluginSpec.hexdoc_book_plugin)
-        for plugin in flatten(caller.try_call()):
+        for plugin in flatten_hook_returns(caller.try_call()):
             self.book_plugins[plugin.modid] = plugin
 
     def _init_mod_plugins(self):
         caller = self._hook_caller(PluginSpec.hexdoc_mod_plugin)
-        for plugin in flatten(
+        for plugin in flatten_hook_returns(
             caller.try_call(
                 branch=self.branch,
                 props=self.props,
             )
         ):
             self.mod_plugins[plugin.modid] = plugin
-            self.item_image_types += flatten([plugin.item_image_types()])
+            self.item_image_types += flatten_hook_return(plugin.item_image_types())
 
     def register(self, plugin: Any, name: str | None = None):
         self.inner.register(plugin, name)
@@ -207,7 +207,7 @@ def validate_format_tree(
     def update_context(self, context: dict[str, Any]) -> Iterator[ValidationContext]:
         caller = self._hook_caller(PluginSpec.hexdoc_update_context)
         if returns := caller.try_call(context=context):
-            yield from flatten(returns)
+            yield from flatten_hook_returns(returns)
 
     def update_jinja_env(self, env: SandboxedEnvironment, modids: Sequence[str]):
         for modid in modids:
@@ -222,7 +222,7 @@ def update_template_args(self, template_args: dict[str, Any]):
 
     def load_resources(self, modid: str) -> Iterator[ModuleType]:
         plugin = self.mod_plugin(modid)
-        for package in flatten([plugin.resource_dirs()]):
+        for package in flatten_hook_return(plugin.resource_dirs()):
             yield import_package(package)
 
     def load_tagged_unions(self) -> Iterator[ModuleType]:
@@ -253,7 +253,7 @@ def _package_loaders_for(self, modids: Iterable[str]):
                         package_name=import_package(package).__name__,
                         package_path=package_path,
                     )
-                    for package, package_path in flatten([result])
+                    for package, package_path in flatten_hook_return(result)
                 ]
             )
 
@@ -289,7 +289,7 @@ def _import_from_hook(
         **kwargs: _P.kwargs,
     ) -> Iterator[ModuleType]:
         packages = self._hook_caller(__spec)(*args, **kwargs)
-        for package in flatten(packages):
+        for package in flatten_hook_returns(packages):
             yield import_package(package)
 
     @overload
@@ -305,12 +305,16 @@ def _hook_caller(self, spec: Callable[_P, _R | None]) -> TypedHookCaller[_P, _R]
         return TypedHookCaller(None, caller)
 
 
-def flatten(values: list[list[_T] | _T] | None) -> Iterator[_T]:
+def flatten_hook_returns(values: HookReturns[_T] | None) -> Iterator[_T]:
     for value in values or []:
-        if isinstance(value, list):
-            yield from value
-        else:
-            yield value
+        yield from flatten_hook_return(value)
+
+
+def flatten_hook_return(values: HookReturn[_T] | None) -> Iterator[_T]:
+    if isinstance(values, list):
+        yield from values
+    else:
+        yield values  # type: ignore
 
 
 def import_package(package: Package) -> ModuleType: