[mypyc] Add librt.threading.Lock class

mypy/typeshed/stubs/librt/librt/threading.pyi

@@ -0,0 +1,16 @@
+from types import TracebackType
+from typing import final
+
+@final
+class Lock:
+    def acquire(self, blocking: bool = True) -> bool: ...
+    def release(self) -> None: ...
+    def locked(self) -> bool: ...
+    def __enter__(self) -> bool: ...
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+        /,
+    ) -> None: ...

mypyc/build.py

@@ -127,6 +127,12 @@ class ModDesc(NamedTuple):
     ),
     ModDesc("librt.time", ["time/librt_time.c"], ["time/librt_time.h"], []),
     ModDesc("librt.random", ["random/librt_random.c"], ["random/librt_random.h"], ["random"]),
+    ModDesc(
+        "librt.threading",
+        ["threading/librt_threading.c"],
+        ["threading/librt_threading.h"],
+        ["threading"],
+    ),
 ]
 
 try:

mypyc/codegen/emitmodule.py

@@ -62,6 +62,7 @@
     LIBRT_BASE64,
     LIBRT_RANDOM,
     LIBRT_STRINGS,
+    LIBRT_THREADING,
     LIBRT_TIME,
     LIBRT_VECS,
     Capsule,
@@ -1261,6 +1262,10 @@ def emit_module_exec_func(
             emitter.emit_line("if (import_librt_time() < 0) {")
             emitter.emit_line("return -1;")
             emitter.emit_line("}")
+        if LIBRT_THREADING in module.dependencies:
+            emitter.emit_line("if (import_librt_threading() < 0) {")
+            emitter.emit_line("return -1;")
+            emitter.emit_line("}")
         if LIBRT_VECS in module.dependencies:
             emitter.emit_line("if (import_librt_vecs() < 0) {")
             emitter.emit_line("return -1;")

mypyc/ir/deps.py

@@ -110,6 +110,7 @@ def get_header(self) -> str:
 LIBRT_VECS: Final = Capsule("librt.vecs")
 LIBRT_TIME: Final = Capsule("librt.time")
 LIBRT_RANDOM: Final = Capsule("librt.random")
+LIBRT_THREADING: Final = Capsule("librt.threading")
 
 BYTES_EXTRA_OPS: Final = SourceDep("bytes_extra_ops.c")
 BYTES_WRITER_EXTRA_OPS: Final = SourceDep("byteswriter_extra_ops.c")

mypyc/ir/rtypes.py

@@ -41,7 +41,7 @@ class to enable the new behavior. In rare cases, adding a new
 from typing import TYPE_CHECKING, ClassVar, Final, Generic, TypeGuard, TypeVar, Union, final
 
 from mypyc.common import HAVE_IMMORTAL, IS_32_BIT_PLATFORM, PLATFORM_SIZE, JsonDict, short_name
-from mypyc.ir.deps import LIBRT_RANDOM, LIBRT_STRINGS, LIBRT_VECS, Dependency
+from mypyc.ir.deps import LIBRT_RANDOM, LIBRT_STRINGS, LIBRT_THREADING, LIBRT_VECS, Dependency
 from mypyc.namegen import NameGenerator
 
 if TYPE_CHECKING:
@@ -550,12 +550,19 @@ def __hash__(self) -> int:
 } | {
     "librt.random.Random": RPrimitive(
         "librt.random.Random", is_unboxed=False, is_refcounted=True, dependencies=(LIBRT_RANDOM,)
-    )
+    ),
+    "librt.threading.Lock": RPrimitive(
+        "librt.threading.Lock",
+        is_unboxed=False,
+        is_refcounted=True,
+        dependencies=(LIBRT_THREADING,),
+    ),
 }
 
 bytes_writer_rprimitive: Final = KNOWN_NATIVE_TYPES["librt.strings.BytesWriter"]
 string_writer_rprimitive: Final = KNOWN_NATIVE_TYPES["librt.strings.StringWriter"]
 random_rprimitive: Final = KNOWN_NATIVE_TYPES["librt.random.Random"]
+lock_rprimitive: Final = KNOWN_NATIVE_TYPES["librt.threading.Lock"]
 
 
 def is_native_rprimitive(rtype: RType) -> bool:

mypyc/irbuild/statement.py

@@ -80,6 +80,7 @@
     c_pyssize_t_rprimitive,
     exc_rtuple,
     is_tagged,
+    lock_rprimitive,
     none_rprimitive,
     object_pointer_rprimitive,
     object_rprimitive,
@@ -115,6 +116,7 @@
     restore_exc_info_op,
 )
 from mypyc.primitives.generic_ops import iter_op, next_raw_op, py_delattr_op
+from mypyc.primitives.librt_threading_ops import lock_acquire_op, lock_release_op
 from mypyc.primitives.misc_ops import (
     check_stop_op,
     coro_op,
@@ -1108,6 +1110,10 @@ def transform_with(
     al = "a" if is_async else ""
 
     mgr_v = builder.accept(expr)
+
+    if not is_async and mgr_v.type == lock_rprimitive:
+        transform_with_lock(builder, mgr_v, target, body, line)
+        return
     is_native = isinstance(mgr_v.type, RInstance)
     if is_native:
         value = builder.add(MethodCall(mgr_v, f"__{al}enter__", args=[], line=line))
@@ -1179,6 +1185,32 @@ def finally_body() -> None:
     )
 
 
+def transform_with_lock(
+    builder: IRBuilder, mgr_v: Value, target: Lvalue | None, body: GenFunc, line: int
+) -> None:
+    """Optimized 'with' for librt.threading.Lock.
+
+    Generate a simple try/finally with direct acquire/release calls.
+    Lock.__exit__ never suppresses exceptions, so we don't need the
+    full PEP 343 try/except/finally machinery.
+    """
+    # __enter__: acquire the lock
+    value = builder.primitive_op(lock_acquire_op, [mgr_v], line)
+
+    mgr = builder.maybe_spill(mgr_v)
+
+    def try_body() -> None:
+        if target:
+            builder.assign(builder.get_assignment_target(target), value, line)
+        body()
+
+    def finally_body() -> None:
+        # __exit__: release the lock (ignoring exception info)
+        builder.primitive_op(lock_release_op, [builder.read(mgr, line)], line)
+
+    transform_try_finally_stmt(builder, try_body, finally_body, line)
+
+
 def transform_with_stmt(builder: IRBuilder, o: WithStmt) -> None:
     # Generate separate logic for each expr in it, left to right
     def generate(i: int) -> None:

mypyc/lib-rt/setup.py

@@ -164,5 +164,11 @@ def run(self) -> None:
                 include_dirs=["."],
                 extra_compile_args=cflags,
             ),
+            Extension(
+                "librt.threading",
+                ["threading/librt_threading.c"],
+                include_dirs=[".", "threading"],
+                extra_compile_args=cflags,
+            ),
         ]
     )