Make accesses to f.func_code return the same object

src/core/types.h

@@ -264,6 +264,7 @@ public:
 
 typedef std::vector<CompiledFunction*> FunctionList;
 struct CallRewriteArgs;
+class BoxedCode;
 class CLFunction {
 public:
     int num_args;
@@ -278,6 +279,9 @@ public:
     CompiledFunction* always_use_version; // if this version is set, always use it (for unboxed cases)
     std::unordered_map<const OSREntryDescriptor*, CompiledFunction*> osr_versions;
 
+    // Please use codeForFunction() to access this:
+    BoxedCode* code_obj;
+
     // Functions can provide an "internal" version, which will get called instead
     // of the normal dispatch through the functionlist.
     // This can be used to implement functions which know how to rewrite themselves,
@@ -289,12 +293,12 @@ public:
     CLFunction(int num_args, int num_defaults, bool takes_varargs, bool takes_kwargs,
                std::unique_ptr<SourceInfo> source)
         : num_args(num_args), num_defaults(num_defaults), takes_varargs(takes_varargs), takes_kwargs(takes_kwargs),
-          source(std::move(source)), param_names(this->source->ast), always_use_version(NULL) {
+          source(std::move(source)), param_names(this->source->ast), always_use_version(NULL), code_obj(NULL) {
         assert(num_args >= num_defaults);
     }
     CLFunction(int num_args, int num_defaults, bool takes_varargs, bool takes_kwargs, const ParamNames& param_names)
         : num_args(num_args), num_defaults(num_defaults), takes_varargs(takes_varargs), takes_kwargs(takes_kwargs),
-          source(nullptr), param_names(param_names), always_use_version(NULL) {
+          source(nullptr), param_names(param_names), always_use_version(NULL), code_obj(NULL) {
         assert(num_args >= num_defaults);
     }
 

src/runtime/code.cpp

@@ -102,12 +102,17 @@ public:
     }
 };
 
-Box* codeForFunction(BoxedFunction* f) {
-    return new BoxedCode(f->f);
+Box* codeForCLFunction(CLFunction* f) {
+    if (!f->code_obj) {
+        f->code_obj = new BoxedCode(f);
+        // CLFunctions don't currently participate in GC.  They actually never get freed currently.
+        gc::registerPermanentRoot(f->code_obj);
+    }
+    return f->code_obj;
 }
 
-Box* codeForCLFunction(CLFunction* f) {
-    return new BoxedCode(f);
+Box* codeForFunction(BoxedFunction* f) {
+    return codeForCLFunction(f->f);
 }
 
 CLFunction* clfunctionFromCode(Box* code) {

src/runtime/types.cpp

@@ -838,7 +838,6 @@ static Box* builtinFunctionOrMethodName(Box* b, void*) {
 static Box* functionCode(Box* self, void*) {
     assert(self->cls == function_cls);
     BoxedFunction* func = static_cast<BoxedFunction*>(self);
-    // This fails "f.func_code is f.func_code"
     return codeForFunction(func);
 }
 
@@ -2424,6 +2423,7 @@ void setupRuntime() {
                            new BoxedFunction(boxRTFunction((void*)functionCall, UNKNOWN, 1, 0, true, true)));
     function_cls->giveAttr("__nonzero__", new BoxedFunction(boxRTFunction((void*)functionNonzero, BOXED_BOOL, 1)));
     function_cls->giveAttr("func_code", new (pyston_getset_cls) BoxedGetsetDescriptor(functionCode, NULL, NULL));
+    function_cls->giveAttr("__code__", function_cls->getattr("func_code"));
     function_cls->giveAttr("func_name", function_cls->getattr("__name__"));
     function_cls->giveAttr("func_defaults",
                            new (pyston_getset_cls) BoxedGetsetDescriptor(functionDefaults, functionSetDefaults, NULL));

test/tests/code_test.py

@@ -17,3 +17,6 @@ def f():
     pass
 print f.func_defaults
 print f.func_code.co_firstlineno
+
+# pytest uses this:
+print f.__code__ is f.__code__