Some logic improvements + attempt to revert on failed import

Author: FasterSpeeding

CHANGELOG.md

@@ -14,8 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - [BasicLocaliser.set_variants][tanjun.dependencies.BasicLocaliser.set_variants]
   now tries to normalise the locale keys of variants passed as keyword arguments
   to match the [hikari.Locale][hikari.locales.Locale] values.
-- The hot-reloader no-longer keeps a module's old components loaded after
-  a top-level error (e.g. syntax error) is raised while re-importing the module.
+- The hot-reloader no-longer reverts to a module's old components loaded after
+  the new module's loaders fail or is found to have no loaders.
 
 ### Fixed
 - Module unloaders will now be called with the correct old module's state

tanjun/clients.py

@@ -2627,12 +2627,8 @@ def load_modules(self, *modules: typing.Union[str, pathlib.Path]) -> Self:
             with _WrapLoadError(errors.FailedModuleImport, module_path):
                 module = load_module()
 
-            try:
+            with _EXPECT_ITER_END:
                 generator.send(module)
-            except StopIteration:
-                pass
-            else:
-                raise RuntimeError("Generator didn't finish")
 
         return self
 
@@ -2648,12 +2644,8 @@ async def load_modules_async(self, *modules: typing.Union[str, pathlib.Path]) ->
             with _WrapLoadError(errors.FailedModuleImport, module_path):
                 module = await loop.run_in_executor(None, load_module)
 
-            try:
+            with _EXPECT_ITER_END:
                 generator.send(module)
-            except StopIteration:
-                pass
-            else:
-                raise RuntimeError("Generator didn't finish")
 
     def unload_modules(self, *modules: typing.Union[str, pathlib.Path]) -> Self:
         # <<inherited docstring from tanjun.ab.Client>>.
@@ -2679,7 +2671,7 @@ def unload_modules(self, *modules: typing.Union[str, pathlib.Path]) -> Self:
 
     def _reload_module(
         self, module_path: typing.Union[str, pathlib.Path], /
-    ) -> collections.Generator[collections.Callable[[], types.ModuleType], types.ModuleType, None]:
+    ) -> collections.Generator[collections.Callable[[], types.ModuleType], types.ModuleType | None, None]:
         if isinstance(module_path, str):
             old_module = self._modules.get(module_path)
             load_module: typing.Optional[_ReloadModule] = None
@@ -2697,25 +2689,33 @@ def _reload_module(
         _LOGGER.info("Reloading %s", module_path)
 
         old_loaders = _get_loaders(old_module, module_path)
-        modules_dict.pop(module_path)
         with _WrapLoadError(errors.FailedModuleUnload, module_path):
-            # This will never raise MissingLoaders as we assert this earlier
             self._call_unloaders(module_path, old_loaders)
 
         module = yield load_module
 
-        loaders = _get_loaders(module, module_path)
+        if not module:
+            # Indicates the newer version of this module couldn't be imported
+            try:
+                self._call_loaders(module_path, old_loaders)
+
+            except Exception as exc:
+                # TODO: exc here is annoyingly already chained with another
+                # error which would've already been logged
+                _LOGGER.debug("Failed to revert %s", module_path, exc_info=exc)
+                modules_dict.pop(module_path)
+
+            return
 
-        # We assert that the new module has loaders early to avoid unnecessarily
-        # unloading then rolling back when we know it's going to fail to load.
-        if not any(loader.has_load for loader in loaders):
-            raise errors.ModuleMissingLoaders(f"Didn't find any loaders in new {module_path}", module_path)
+        loaders = _get_loaders(module, module_path)
 
         try:
-            # This will never raise MissingLoaders as we assert this earlier
             self._call_loaders(module_path, loaders)
+        except errors.ModuleMissingLoaders:
+            modules_dict.pop(module_path)
+            raise
         except Exception as exc:
-            self._call_loaders(module_path, old_loaders)
+            modules_dict.pop(module_path)
             raise errors.FailedModuleLoad(module_path) from exc
         else:
             modules_dict[module_path] = module
@@ -2728,15 +2728,18 @@ def reload_modules(self, *modules: typing.Union[str, pathlib.Path]) -> Self:
 
             generator = self._reload_module(module_path)
             load_module = next(generator)
-            with _WrapLoadError(errors.FailedModuleLoad, module_path):
-                module = load_module()
 
             try:
+                module = load_module()
+
+            except Exception as exc:
+                with _EXPECT_ITER_END:
+                    generator.send(None)
+
+                raise errors.FailedModuleLoad(module_path) from exc
+
+            with _EXPECT_ITER_END:
                 generator.send(module)
-            except StopIteration:
-                pass
-            else:
-                raise RuntimeError("Generator didn't finish")
 
         return self
 
@@ -2749,17 +2752,18 @@ async def reload_modules_async(self, *modules: typing.Union[str, pathlib.Path])
 
             generator = self._reload_module(module_path)
             load_module = next(generator)
-            with _WrapLoadError(errors.FailedModuleLoad, module_path):
-                module = await loop.run_in_executor(None, load_module)
 
             try:
-                generator.send(module)
+                module = await loop.run_in_executor(None, load_module)
 
-            except StopIteration:
-                pass
+            except Exception as exc:
+                with _EXPECT_ITER_END:
+                    generator.send(None)
 
-            else:
-                raise RuntimeError("Generator didn't finish")
+                raise errors.FailedModuleLoad(module_path) from exc
+
+            with _EXPECT_ITER_END:
+                generator.send(module)
 
     def set_type_dependency(self, type_: type[_T], value: _T, /) -> Self:
         # <<inherited docstring from tanjun.abc.Client>>.
@@ -3185,6 +3189,27 @@ def __exit__(
             raise self._error(*self._args, **self._kwargs) from exc
 
 
+class _ExpectIterEnd:
+    __slots__ = ()
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(
+        self,
+        exc_type: typing.Optional[type[BaseException]],
+        exc: typing.Optional[BaseException],
+        exc_tb: typing.Optional[types.TracebackType],
+    ) -> bool:
+        if exc_type is None:
+            raise RuntimeError("Generator didn't finish") from None
+
+        return exc_type is StopIteration
+
+
+_EXPECT_ITER_END = _ExpectIterEnd()
+
+
 def _try_deregister_listener(
     interaction_server: hikari.api.InteractionServer,
     interaction_type: typing.Any,