Update specialize.py

Author: BobTheBuidler

mypyc/irbuild/specialize.py

@@ -96,7 +96,10 @@
     str_encode_utf8_strict,
 )
 from mypyc.primitives.tuple_ops import new_tuple_set_item_op
-from mypyc.primitives.weakref_ops import new_proxy_op
+from mypyc.primitives.weakref_ops import (
+    new_proxy_op,
+    new_ref_op,
+)
 
 # Specializers are attempted before compiling the arguments to the
 # function.  Specializers can return None to indicate that they failed
@@ -891,29 +894,51 @@ def translate_ord(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value
     return None
 
 
-@specialize_function("weakref.proxy")
+@specialize_function("weakref.ReferenceType")
 def translate_proxy(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
     """
-    Special case for weakref.proxy(...).
+    Special case for weakref.ref(...).
 
     If there's only one argument (the object), pass None for the callback.
     If there are two arguments, use them as (object, callback).
     Otherwise, fallback to a normal Python call by returning None.
     """
-    line = expr.line
+    obj_val = builder.accept(expr.args[0])
     num_args = len(expr.args)
-
-    # If the user supplied something other than 1 or 2 args, just fail
-    # and let MyPyC do a generic Python call.
-    if num_args not in (1, 2):
+    if num_args == 1:
+        callback_val = builder.none_object()
+    elif num_args == 2:
+        callback_val = builder.accept(expr.args[1])
+    else:
+        # If the user supplied something other than 1 or 2 args, just fail
+        # and let MyPyC do a generic Python call.
         # NOTE: should we raise a compiler error here?
         return None
 
+    # Emit a direct C call to PyWeakref_NewProxy
+    return builder.call_c(new_ref_op, [obj_val, callback_val], expr.line)
+
+
+@specialize_function("weakref.proxy")
+def translate_proxy(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """
+    Special case for weakref.proxy(...).
+
+    If there's only one argument (the object), pass None for the callback.
+    If there are two arguments, use them as (object, callback).
+    Otherwise, fallback to a normal Python call by returning None.
+    """
     obj_val = builder.accept(expr.args[0])
+    num_args = len(expr.args)
     if num_args == 1:
         callback_val = builder.none_object()
-    else:
+    elif num_args == 2:
         callback_val = builder.accept(expr.args[1])
+    else:
+        # If the user supplied something other than 1 or 2 args, just fail
+        # and let MyPyC do a generic Python call.
+        # NOTE: should we raise a compiler error here?
+        return None
 
     # Emit a direct C call to PyWeakref_NewProxy
-    return builder.call_c(new_proxy_op, [obj_val, callback_val], line)
+    return builder.call_c(new_proxy_op, [obj_val, callback_val], expr.line)