From 35e87e9bbf3366c8ce5a6d8988daf93f0a78fb95 Mon Sep 17 00:00:00 2001
From: Alessio Bogon <778703+youtux@users.noreply.github.com>
Date: Mon, 2 Sep 2024 17:20:47 +0200
Subject: [PATCH] Add misc missing types (#163)

This PR adds a few missing types:
* `ColorFormatter`
* `get_current_task`

and it fixes the definition of `Context.headers` to be a dict (mutable)
---
 celery-stubs/_state.pyi    |  2 ++
 celery-stubs/app/task.pyi  |  2 +-
 celery-stubs/utils/log.pyi | 12 +++++++++++-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/celery-stubs/_state.pyi b/celery-stubs/_state.pyi
index 4bcb4d2..0d02650 100644
--- a/celery-stubs/_state.pyi
+++ b/celery-stubs/_state.pyi
@@ -5,3 +5,5 @@ from celery.app.task import Task
 
 current_app: Celery
 current_task: Task[Any, Any]
+
+def get_current_task() -> Task[Any, Any]: ...
diff --git a/celery-stubs/app/task.pyi b/celery-stubs/app/task.pyi
index 0733944..a65910a 100644
--- a/celery-stubs/app/task.pyi
+++ b/celery-stubs/app/task.pyi
@@ -37,7 +37,7 @@ class Context:
     expires: int | None
     group: str | None
     group_index: int | None
-    headers: Mapping[str, Any] | None
+    headers: dict[str, Any] | None
     hostname: str | None
     id: str | None
     ignore_result: bool
diff --git a/celery-stubs/utils/log.pyi b/celery-stubs/utils/log.pyi
index 3b2d594..689feef 100644
--- a/celery-stubs/utils/log.pyi
+++ b/celery-stubs/utils/log.pyi
@@ -1,4 +1,7 @@
-from logging import Logger
+from logging import Formatter, Logger
+from typing import ClassVar
+
+from celery.utils.term import colored
 
 class LoggingProxy: ...
 
@@ -7,3 +10,10 @@ def get_task_logger(name: str) -> Logger: ...
 
 task_logger: Logger
 worker_logger: Logger
+
+class ColorFormatter(Formatter):
+    use_color: bool
+
+    COLORS: ClassVar[dict[str, colored]]
+    colors: ClassVar[dict[str, colored]]
+    def __init__(self, fmt: str | None = ..., use_color: bool = ...) -> None: ...