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/iot/iotdevice.py

@@ -19,13 +19,14 @@ import functools
 import inspect
 import logging
 from datetime import datetime, timedelta
-from typing import Any, Mapping, Sequence, cast
+from typing import Any, Mapping, Sequence, cast, overload
 
 from ..device import Device, WifiNetwork
 from ..deviceconfig import DeviceConfig
 from ..emeterstatus import EmeterStatus
 from ..exceptions import KasaException
 from ..feature import Feature
+from ..module import ModuleT
 from ..protocol import BaseProtocol
 from .iotmodule import IotModule
 from .modules import Emeter, Time
@@ -201,6 +202,26 @@ class IotDevice(Device):
         """Return the device modules."""
         return self._modules
 
+    @overload
+    def get_module(self, module_type: type[ModuleT]) -> ModuleT | None: ...
+
+    @overload
+    def get_module(self, module_type: str) -> IotModule | None: ...
+
+    def get_module(
+        self, module_type: type[ModuleT] | str
+    ) -> ModuleT | IotModule | None:
+        """Return the module from the device modules or None if not present."""
+        if isinstance(module_type, str):
+            module_name = module_type.lower()
+        elif issubclass(module_type, IotModule):
+            module_name = module_type.__name__.lower()
+        else:
+            return None
+        if module_name in self.modules:
+            return self.modules[module_name]
+        return None
+
     def add_module(self, name: str, module: IotModule):
         """Register a module."""
         if name in self.modules: