Make get_module return typed module (#892)

Passing in a string still works and returns either `IotModule` or
`SmartModule` type when called on `IotDevice` or `SmartDevice`
respectively. When calling on `Device` will return `Module` type.

Passing in a module type is then typed to that module, i.e.:
```py
smartdev.get_module(FanModule)  # type is FanModule
smartdev.get_module("FanModule")  # type is SmartModule
```
Only thing this doesn't do is check that you can't pass an `IotModule`
to a `SmartDevice.get_module()`. However there is a runtime check which
will return null if the passed `ModuleType` is not a subclass of
`SmartModule`.

Many thanks to @cdce8p for helping with this.

kasa/smart/smartdevice.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import base64
 import logging
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Mapping, Sequence, cast
+from typing import Any, Mapping, Sequence, cast, overload
 
 from ..aestransport import AesTransport
 from ..bulb import HSV, Bulb, BulbPreset, ColorTempRange
@@ -16,6 +16,7 @@ from ..emeterstatus import EmeterStatus
 from ..exceptions import AuthenticationError, DeviceError, KasaException, SmartErrorCode
 from ..fan import Fan
 from ..feature import Feature
+from ..module import ModuleT
 from ..smartprotocol import SmartProtocol
 from .modules import (
     Brightness,
@@ -28,11 +29,10 @@ from .modules import (
     Firmware,
     TimeModule,
 )
+from .smartmodule import SmartModule
 
 _LOGGER = logging.getLogger(__name__)
 
-if TYPE_CHECKING:
-    from .smartmodule import SmartModule
 
 # List of modules that wall switches with children, i.e. ks240 report on
 # the child but only work on the parent.  See longer note below in _initialize_modules.
@@ -305,8 +305,22 @@ class SmartDevice(Bulb, Fan, Device):
             for feat in module._module_features.values():
                 self._add_feature(feat)
 
-    def get_module(self, module_name) -> SmartModule | None:
+    @overload
+    def get_module(self, module_type: type[ModuleT]) -> ModuleT | None: ...
+
+    @overload
+    def get_module(self, module_type: str) -> SmartModule | None: ...
+
+    def get_module(
+        self, module_type: type[ModuleT] | str
+    ) -> ModuleT | SmartModule | None:
         """Return the module from the device modules or None if not present."""
+        if isinstance(module_type, str):
+            module_name = module_type
+        elif issubclass(module_type, SmartModule):
+            module_name = module_type.__name__
+        else:
+            return None
         if module_name in self.modules:
             return self.modules[module_name]
         elif self._exposes_child_modules: