add type hints

Author: vlad-scherbich

tests/profiling/collector/test_asyncio.py

@@ -51,41 +51,30 @@ def teardown_method(self):
             except Exception as e:
                 print("Error while deleting file: ", e)
 
-    def test_subclassing_wrapped_lock(self):
-        """Test that subclassing asyncio.Lock works when profiling is active.
+    def test_subclassing_wrapped_lock(self) -> None:
+        """Test that subclassing of a wrapped lock type when profiling is active."""
+        from typing import Optional
 
-        Regression test for a bug reported by neo4j users where their library defines:
-            class AsyncRLock(asyncio.Lock): ...
-
-        With profiling enabled, this would fail with:
-            TypeError: _LockAllocatorWrapper.__init__() takes 2 positional arguments but 4 were given
-
-        The fix implements __mro_entries__ on _LockAllocatorWrapper (PEP 560).
-        """
         from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
 
         with collector_asyncio.AsyncioLockCollector(capture_pct=100):
-            # Verify the wrapper is in place
             assert isinstance(asyncio.Lock, _LockAllocatorWrapper)
 
             # This should NOT raise TypeError
-            class AsyncRLock(asyncio.Lock):
-                """A custom lock that extends asyncio.Lock (like neo4j does)."""
-
+            class CustomLock(asyncio.Lock):  # type: ignore[misc]
                 def __init__(self) -> None:
                     super().__init__()
-                    self._owner = None
-                    self._count = 0
+                    self._owner: Optional[int] = None
+                    self._count: int = 0
 
-            # Verify the subclass can be instantiated
-            custom_lock = AsyncRLock()
+            # Verify subclassing and functionality
+            custom_lock: CustomLock = CustomLock()
             assert hasattr(custom_lock, "_owner")
             assert hasattr(custom_lock, "_count")
             assert custom_lock._owner is None
             assert custom_lock._count == 0
 
-            # Verify async lock functionality still works
-            async def test_async_lock():
+            async def test_async_lock() -> None:
                 await custom_lock.acquire()
                 assert custom_lock.locked()
                 custom_lock.release()

tests/profiling/collector/test_threading.py

@@ -565,34 +565,23 @@ def __init__(self, lock_class: LockTypeClass) -> None:
             Foobar(self.lock_class)
 
     def test_subclassing_wrapped_lock(self) -> None:
-        """Test that subclassing a wrapped lock type works correctly.
-
-        Regression test for a bug where third-party libraries (e.g. neo4j) that subclass
-        asyncio.Lock would fail with:
-            TypeError: _LockAllocatorWrapper.__init__() takes 2 positional arguments but 4 were given
-
-        The fix implements __mro_entries__ on _LockAllocatorWrapper (PEP 560).
-        """
+        """Test that subclassing of a wrapped lock type when profiling is active."""
         from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
 
         with self.collector_class():
-            # Verify the wrapper is in place
             assert isinstance(self.lock_class, _LockAllocatorWrapper)
 
             # This should NOT raise TypeError
             class CustomLock(self.lock_class):  # type: ignore[misc]
-                """A custom lock that extends the wrapped lock type."""
-
                 def __init__(self) -> None:
                     super().__init__()
-                    self.custom_attr = "test"
+                    self.custom_attr: str = "test"
 
-            # Verify the subclass works correctly
-            custom_lock = CustomLock()
+            # Verify subclassing and functionality
+            custom_lock: CustomLock = CustomLock()
             assert hasattr(custom_lock, "custom_attr")
             assert custom_lock.custom_attr == "test"
 
-            # Verify lock functionality still works
             assert custom_lock.acquire()
             custom_lock.release()