Improve types for the SQS store (#13684)

giograno · web-flow · commit 2a499fde7a28 · 2026-02-03T17:15:03.000+01:00
diff --git a/localstack-core/localstack/services/sqs/models.py b/localstack-core/localstack/services/sqs/models.py
@@ -56,7 +56,7 @@
 class SqsMessage:
     message: Message
     created: float
-    visibility_timeout: int
+    visibility_timeout: int | None
     receive_count: int
     delay_seconds: int | None
     receipt_handles: set[str]
@@ -65,9 +65,7 @@ class SqsMessage:
     visibility_deadline: float | None
     deleted: bool
     priority: float
-    message_deduplication_id: str
-    message_group_id: str
-    sequence_number: str
+    sequence_number: str | None
 
     def __init__(
         self,
@@ -85,6 +83,7 @@ def __init__(
         self.delay_seconds = None
         self.last_received = None
         self.first_received = None
+        self.visibility_timeout = None
         self.visibility_deadline = None
         self.deleted = False
         self.priority = priority
@@ -270,28 +269,36 @@ class MessageMoveTask:
     # configurable fields
     source_arn: str
     """The arn of the DLQ the messages are currently in."""
-    destination_arn: str | None = None
+    destination_arn: str
     """If the DestinationArn is not specified, the original source arn will be used as target."""
-    max_number_of_messages_per_second: int | None = None
+    max_number_of_messages_per_second: int | None
 
     # dynamic fields
     task_id: str
-    status: str = MessageMoveTaskStatus.CREATED
-    started_timestamp: datetime | None = None
-    approximate_number_of_messages_moved: int | None = None
-    approximate_number_of_messages_to_move: int | None = None
-    failure_reason: str | None = None
+    status: str
+    started_timestamp: datetime | None
+    approximate_number_of_messages_moved: int | None
+    approximate_number_of_messages_to_move: int | None
+    failure_reason: str | None
 
     cancel_event: threading.Event
 
     def __init__(
-        self, source_arn: str, destination_arn: str, max_number_of_messages_per_second: int = None
+        self,
+        source_arn: str,
+        destination_arn: str,
+        max_number_of_messages_per_second: int | None = None,
     ):
         self.task_id = long_uid()
         self.source_arn = source_arn
         self.destination_arn = destination_arn
         self.max_number_of_messages_per_second = max_number_of_messages_per_second
         self.cancel_event = threading.Event()
+        self.status = MessageMoveTaskStatus.CREATED
+        self.started_timestamp = None
+        self.approximate_number_of_messages_moved = None
+        self.approximate_number_of_messages_to_move = None
+        self.failure_reason = None
 
     def mark_started(self):
         self.started_timestamp = datetime.utcnow()
@@ -318,6 +325,7 @@ class SqsQueue:
     # Simulating an ordered set in python. Only the keys are used and of interest.
     inflight: dict[SqsMessage, None]
     receipts: dict[str, SqsMessage]
+    mutex: threading.RLock
 
     def __init__(self, name: str, region: str, account_id: str, attributes=None, tags=None) -> None:
         self.name = name
@@ -1374,7 +1382,7 @@ def clear(self):
 
 
 class SqsStore(BaseStore):
-    queues: dict[str, SqsQueue] = LocalAttribute(default=dict)
+    queues: dict[str, FifoQueue | StandardQueue] = LocalAttribute(default=dict)
 
     deleted: dict[str, float] = LocalAttribute(default=dict)
 
diff --git a/localstack-core/localstack/services/sqs/queue.py b/localstack-core/localstack/services/sqs/queue.py
@@ -1,16 +1,19 @@
 import time
 from queue import Empty, PriorityQueue, Queue
+from typing import Generic, TypeVar
 
+T = TypeVar("T")
 
-class InterruptibleQueue(Queue):
+
+class InterruptibleQueue(Queue, Generic[T]):
     # is_shutdown is used to check whether we have triggered a shutdown of the Queue
     is_shutdown: bool
 
-    def __init__(self, maxsize=0):
+    def __init__(self, maxsize: int = 0):
         super().__init__(maxsize)
         self.is_shutdown = False
 
-    def get(self, block=True, timeout=None):
+    def get(self, block: bool = True, timeout: float | None = None) -> T:
         with self.not_empty:
             if self.is_shutdown:
                 raise Empty
@@ -35,7 +38,7 @@ def get(self, block=True, timeout=None):
             self.not_full.notify()
             return item
 
-    def shutdown(self):
+    def shutdown(self) -> None:
         """
         `shutdown` signals to stop all current and future `Queue.get` calls from executing.
 
@@ -46,5 +49,5 @@ def shutdown(self):
             self.not_empty.notify_all()
 
 
-class InterruptiblePriorityQueue(PriorityQueue, InterruptibleQueue):
+class InterruptiblePriorityQueue(PriorityQueue, InterruptibleQueue[T], Generic[T]):
     pass
diff --git a/localstack-core/localstack/testing/pytest/container.py b/localstack-core/localstack/testing/pytest/container.py
@@ -23,7 +23,7 @@
 LOG = logging.getLogger(__name__)
 
 ENV_TEST_CONTAINER_MOUNT_SOURCES = "TEST_CONTAINER_MOUNT_SOURCES"
-"""Environment variable used to indicate that we should mount LocalStack  source files into the container."""
+"""Environment variable used to indicate that we should mount LocalStack source files into the container."""
 
 ENV_TEST_CONTAINER_MOUNT_DEPENDENCIES = "TEST_CONTAINER_MOUNT_DEPENDENCIES"
 """Environment variable used to indicate that we should mount dependencies into the container."""
