Improved and fixed some types

Author: viblo

pymunk/_weakkeysview.py

@@ -17,9 +17,8 @@ def __iter__(self) -> Iterator[KT]:
     def __len__(self) -> int:
         return len(self._weak_dict)
 
-    def __contains__(self, key: KT) -> bool:
+    def __contains__(self, key: object) -> bool:
         return key in self._weak_dict
 
     def __repr__(self) -> str:
-        # Provide a string representation of the keys view
         return f"{self.__class__.__name__}({list(self._weak_dict.keys())})"

pymunk/body.py

@@ -114,7 +114,7 @@ class Body(PickleMixin, TypingAttrMixing, object):
     _position_func: Optional[_PositionFunc] = None
     _velocity_func: Optional[_VelocityFunc] = None
 
-    _dead_ref: weakref.ref = weakref.ref(set())
+    _dead_ref: weakref.ref[Any] = weakref.ref(set())
 
     def __init__(
         self, mass: float = 0, moment: float = 0, body_type: _BodyType = DYNAMIC
@@ -220,9 +220,9 @@ def freebody(cp_body: ffi.CData) -> None:
         elif body_type == Body.STATIC:
             self._body = ffi.gc(lib.cpBodyNewStatic(), freebody)
 
-        self._space: weakref.ref = Body._dead_ref
+        self._space: weakref.ref["Space"] = Body._dead_ref
 
-        self._constraints: WeakKeyDictionary = WeakKeyDictionary()
+        self._constraints: WeakKeyDictionary["Constraint", None] = WeakKeyDictionary()
         self._shapes: dict["Shape", None] = {}
 
         d = ffi.new_handle(self)

pymunk/shapes.py

@@ -1,7 +1,7 @@
 __docformat__ = "reStructuredText"
 
 import weakref
-from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple
 
 if TYPE_CHECKING:
     from .body import Body
@@ -43,11 +43,12 @@ class Shape(PickleMixin, TypingAttrMixing, object):
     _pickle_attrs_skip = PickleMixin._pickle_attrs_skip + ["mass", "density"]
 
     _space = None  # Weak ref to the space holding this body (if any)
-    _dead_ref = weakref.ref(set())
+    _dead_ref: weakref.ref[Any] = weakref.ref(set())
 
     def __init__(self, shape: "Shape") -> None:
         self._shape = shape
-        self._body: weakref.ref = weakref.ref(shape.body)
+        assert shape.body != None
+        self._body: weakref.ref["Body"] = weakref.ref(shape.body)
 
     def _init(self, body: Optional["Body"], _shape: ffi.CData) -> None:
 

pymunk/space.py

@@ -1,5 +1,6 @@
 __docformat__ = "reStructuredText"
 
+import math
 import platform
 import weakref
 from typing import (

pymunk/tests/test_arbiter.py

@@ -197,9 +197,8 @@ def testTotalKE(self) -> None:
 
         s.add(b1, c1, b2, c2)
 
-        def post_solve(arb: p.Arbiter, space: p.Space, data: Any) -> bool:
+        def post_solve(arb: p.Arbiter, space: p.Space, data: Any) -> None:
             self.assertAlmostEqual(arb.total_ke, 43.438914027)
-            return True
 
         s.add_collision_handler(1, 2).post_solve = post_solve
 
@@ -292,10 +291,9 @@ def separate1(arb: p.Arbiter, space: p.Space, data: Any) -> None:
 
         self.called2 = False
 
-        def separate2(arb: p.Arbiter, space: p.Space, data: Any) -> bool:
+        def separate2(arb: p.Arbiter, space: p.Space, data: Any) -> None:
             self.called2 = True
             self.assertTrue(arb.is_removal)
-            return True
 
         s.add_collision_handler(1, 2).separate = separate2
         s.remove(b1, c1)

pymunk/tests/test_common.py

@@ -177,7 +177,7 @@ def separate(arbiter: p.Arbiter, space: p.Space, data: Any) -> None:
     def testWeakKeysView(self) -> None:
         x1, x2, x3 = p.Body(1), p.Body(2), p.Body(3)
 
-        d = weakref.WeakKeyDictionary()
+        d: weakref.WeakKeyDictionary[p.Body, int] = weakref.WeakKeyDictionary()
 
         d[x1] = 1
         d[x2] = 2