Add modern type hints. (#202)

* Use `typing-extensions` to support older versions. This is needed for `ParamSpec` and `TypeAlias`.
Mulugruntz · Apr 2, 2023 · 42676c1 · 42676c1
1 parent 87801fa
commit 42676c1
Show file tree

Hide file tree

Showing 9 changed files with 187 additions and 120 deletions.
diff --git a/README.md b/README.md
@@ -1,4 +1,4 @@
-# celery-pubsub 2.0.0-beta2
+# celery-pubsub 2.0.0-beta3
 
 
 [![Build and Test](https://github.com/Mulugruntz/celery-pubsub/actions/workflows/build.yml/badge.svg)](https://github.com/Mulugruntz/celery-pubsub/actions/workflows/build.yml)
@@ -105,6 +105,7 @@ celery_pubsub.publish('some.very.good.test', 42)  # task 3 only
     * Drop support for Pypy 2.7 and 3.6.
     * Add support for Pypy 3.8 and 3.9.
     * Add support for CPython 3.11.
+    * Type hints are now directly in the code. No more stubs files.
 * 1.0.2
     * Add stubs file for type hinting.
 * 1.0.1
@@ -190,30 +191,30 @@ celery_pubsub.publish('some.very.good.test', 42)  # task 3 only
 [badge-m_linux_pypy3.9_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/m_linux_pypy-3.9_celery5/shields
 
 [//]: # (Status in tagged version)
-[badge-t_linux_3.7_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.7_celery3/shields
-[badge-t_linux_3.7_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.7_celery4/shields
-[badge-t_linux_3.7_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.7_celery5/shields
+[badge-t_linux_3.7_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.7_celery3/shields
+[badge-t_linux_3.7_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.7_celery4/shields
+[badge-t_linux_3.7_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.7_celery5/shields
 
-[badge-t_linux_3.8_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.8_celery3/shields
-[badge-t_linux_3.8_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.8_celery4/shields
-[badge-t_linux_3.8_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.8_celery5/shields
+[badge-t_linux_3.8_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.8_celery3/shields
+[badge-t_linux_3.8_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.8_celery4/shields
+[badge-t_linux_3.8_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.8_celery5/shields
 
-[badge-t_linux_3.9_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.9_celery3/shields
-[badge-t_linux_3.9_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.9_celery4/shields
-[badge-t_linux_3.9_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.9_celery5/shields
+[badge-t_linux_3.9_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.9_celery3/shields
+[badge-t_linux_3.9_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.9_celery4/shields
+[badge-t_linux_3.9_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.9_celery5/shields
 
-[badge-t_linux_3.10_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.10_celery3/shields
-[badge-t_linux_3.10_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.10_celery4/shields
-[badge-t_linux_3.10_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.10_celery5/shields
+[badge-t_linux_3.10_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.10_celery3/shields
+[badge-t_linux_3.10_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.10_celery4/shields
+[badge-t_linux_3.10_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.10_celery5/shields
 
-[badge-t_linux_3.11_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.11_celery3/shields
-[badge-t_linux_3.11_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.11_celery4/shields
-[badge-t_linux_3.11_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_3.11_celery5/shields
+[badge-t_linux_3.11_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.11_celery3/shields
+[badge-t_linux_3.11_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.11_celery4/shields
+[badge-t_linux_3.11_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_3.11_celery5/shields
 
-[badge-t_linux_pypy3.8_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_pypy-3.8_celery3/shields
-[badge-t_linux_pypy3.8_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_pypy-3.8_celery4/shields
-[badge-t_linux_pypy3.8_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_pypy-3.8_celery5/shields
+[badge-t_linux_pypy3.8_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_pypy-3.8_celery3/shields
+[badge-t_linux_pypy3.8_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_pypy-3.8_celery4/shields
+[badge-t_linux_pypy3.8_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_pypy-3.8_celery5/shields
 
-[badge-t_linux_pypy3.9_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_pypy-3.9_celery3/shields
-[badge-t_linux_pypy3.9_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_pypy-3.9_celery4/shields
-[badge-t_linux_pypy3.9_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta2_linux_pypy-3.9_celery5/shields
+[badge-t_linux_pypy3.9_celery3]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_pypy-3.9_celery3/shields
+[badge-t_linux_pypy3.9_celery4]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_pypy-3.9_celery4/shields
+[badge-t_linux_pypy3.9_celery5]: https://byob.yarr.is/Mulugruntz/celery-pubsub/2.0.0-beta3_linux_pypy-3.9_celery5/shields
diff --git a/celery_pubsub/__init__.pyi b/celery_pubsub/__init__.pyi
diff --git a/celery_pubsub/pubsub.py b/celery_pubsub/pubsub.py
@@ -1,3 +1,20 @@
+"""Contains the pubsub manager and the pubsub functions."""
+
+from __future__ import annotations
+
+import typing
+
+if typing.TYPE_CHECKING:  # pragma: no cover
+    from typing_extensions import TypeAlias
+else:
+    try:
+        from typing import TypeAlias as TypeAlias
+    except ImportError:
+        try:
+            from typing_extensions import TypeAlias as TypeAlias
+        except ImportError:
+            TypeAlias = None
+
 import celery
 import re
 
@@ -8,68 +25,76 @@
     "unsubscribe",
 ]
 
+from celery import Task, group
+from celery.result import AsyncResult, EagerResult
+
+
+PA: TypeAlias = typing.Any  # ParamSpec args
+PK: TypeAlias = typing.Any  # ParamSpec kwargs
+P: TypeAlias = typing.Any  # ParamSpec
+R: TypeAlias = typing.Any  # Return type
+
 
-class PubSubManager(object):
-    def __init__(self):
+class PubSubManager:
+    def __init__(self) -> None:
         super(PubSubManager, self).__init__()
-        self.subscribed = set()
-        self.jobs = {}
+        self.subscribed: set[tuple[str, re.Pattern[str], Task[P, R]]] = set()
+        self.jobs: dict[str, group] = {}
 
-    def publish(self, topic, *args, **kwargs):
+    def publish(self, topic: str, *args: PA, **kwargs: PK) -> AsyncResult[R]:
         result = self.get_jobs(topic).delay(*args, **kwargs)
         return result
 
-    def publish_now(self, topic, *args, **kwargs):
-        result = self.get_jobs(topic).apply(args=args, kwargs=kwargs)
+    def publish_now(self, topic: str, *args: PA, **kwargs: PK) -> EagerResult[R]:
+        # Ignoring type because of this: https://github.com/sbdchd/celery-types/issues/111
+        result = self.get_jobs(topic).apply(args=args, kwargs=kwargs)  # type: ignore
         return result
 
-    def subscribe(self, topic, task):
+    def subscribe(self, topic: str, task: Task[P, R]) -> None:
         key = (topic, self._topic_to_re(topic), task)
         if key not in self.subscribed:
             self.subscribed.add(key)
             self.jobs = {}
 
-    def unsubscribe(self, topic, task):
+    def unsubscribe(self, topic: str, task: Task[P, R]) -> None:
         key = (topic, self._topic_to_re(topic), task)
         if key in self.subscribed:
             self.subscribed.discard(key)
             self.jobs = {}
 
-    def get_jobs(self, topic):
+    def get_jobs(self, topic: str) -> group:
         if topic not in self.jobs:
             self._gen_jobs(topic)
         return self.jobs[topic]
 
-    def _gen_jobs(self, topic):
+    def _gen_jobs(self, topic: str) -> None:
         jobs = []
         for job in self.subscribed:
             if job[1].match(topic):
                 jobs.append(job[2].s())
         self.jobs[topic] = celery.group(jobs)
 
     @staticmethod
-    def _topic_to_re(topic):
+    def _topic_to_re(topic: str) -> re.Pattern[str]:
         assert isinstance(topic, str)
         re_topic = topic.replace(".", r"\.").replace("*", r"[^.]+").replace("#", r".+")
         return re.compile(r"^{}$".format(re_topic))
 
 
-_pubsub_manager = None
-if _pubsub_manager is None:  # pragma: no cover
-    _pubsub_manager = PubSubManager()
+_pubsub_manager: PubSubManager = PubSubManager()
 
 
-def publish(topic, *args, **kwargs):
+def publish(topic: str, *args: PA, **kwargs: PK) -> AsyncResult[R]:
     return _pubsub_manager.publish(topic, *args, **kwargs)
 
 
-def publish_now(topic, *args, **kwargs):
+def publish_now(topic: str, *args: PA, **kwargs: PK) -> EagerResult[R]:
     return _pubsub_manager.publish_now(topic, *args, **kwargs)
 
 
-def subscribe(topic, task):
+def subscribe(topic: str, task: Task[P, R]) -> None:
     return _pubsub_manager.subscribe(topic, task)
 
 
-def unsubscribe(topic, task):
+def unsubscribe(topic: str, task: Task[P, R]) -> None:
     return _pubsub_manager.unsubscribe(topic, task)
diff --git a/celery_pubsub/pubsub.pyi b/celery_pubsub/pubsub.pyi
diff --git a/mypy.ini b/mypy.ini
@@ -1,4 +1,4 @@
 [mypy]
 strict = True
 # Only check files in the `celery_pubsub` package
-files = celery_pubsub
+files = celery_pubsub, tests
diff --git a/requirements_test.txt b/requirements_test.txt
@@ -1,6 +1,9 @@
 pytest
 coverage==5.5
+black==23.3.0
+typing-extensions==4.5.0
 
 # Only for CPython (exclude pypy).
 celery-types==0.14.0; implementation_name != 'pypy'
 mypy==1.0.0; implementation_name != 'pypy'
+types-setuptools==67.6.0.6; implementation_name != 'pypy'
diff --git a/setup.py b/setup.py
@@ -34,15 +34,15 @@ def tests_require():
 setuptools.setup(
     name="celery-pubsub",
     packages=["celery_pubsub"],
-    version="2.0.0-beta2",
+    version="2.0.0-beta3",
     description="A Publish and Subscribe library for Celery",
     long_description=long_description(),
     long_description_content_type="text/markdown",
     author="Samuel GIFFARD",
     author_email="[email protected]",
     license="MIT",
     url="https://github.com/Mulugruntz/celery-pubsub",
-    download_url="https://github.com/Mulugruntz/celery-pubsub/tarball/2.0.0-beta2",
+    download_url="https://github.com/Mulugruntz/celery-pubsub/tarball/2.0.0-beta3",
     keywords=["celery", "publish", "subscribe", "pubsub"],
     classifiers=[
         "Development Status :: 5 - Production/Stable",