fix: Load globals that are subclasses or instances of a class being compiled lazily (#67)

Previously, the entire global dict was loaded whenever a function was
compiled. However, this causes issues if there is a global that is an
instance or subclass of a class being compiled, and the class being
compiled reference that global.

Now, when a class is being compiled, instances and subclasses of that
class is excluded and loaded at runtime.

Author: Christopher-Chianelli

jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/CPythonBackedPythonInterpreter.java

@@ -46,6 +46,7 @@ public class CPythonBackedPythonInterpreter implements PythonInterpreter {
     public static BiFunction<OpaquePythonReference, String, PythonLikeObject> lookupAttributeOnPythonReferencePythonFunction;
     public static BiFunction<OpaquePythonReference, String, OpaquePythonReference> lookupPointerForAttributeOnPythonReferencePythonFunction;
     public static BiFunction<OpaquePythonReference, String, OpaquePythonReference[]> lookupPointerArrayForAttributeOnPythonReferencePythonFunction;
+    public static BiConsumer<Map<String, PythonLikeObject>, String> loadObjectFromPythonGlobalDict;
 
     public static TriFunction<OpaquePythonReference, String, Map<Number, PythonLikeObject>, PythonLikeObject> lookupAttributeOnPythonReferenceWithMapPythonFunction;
     public static QuadConsumer<OpaquePythonReference, OpaquePythonReference, String, Object> setAttributeOnPythonReferencePythonFunction;
@@ -154,6 +155,11 @@ public void setPythonReference(PythonLikeObject instance, OpaquePythonReference
 
     @Override
     public PythonLikeObject getGlobal(Map<String, PythonLikeObject> globalsMap, String name) {
+        if (!globalsMap.containsKey(name)) {
+            // This will put 'null' in the map if it doesn't exist, so we don't
+            // do an expensive CPython lookup everytime we are getting an attribute
+            loadObjectFromPythonGlobalDict.accept(globalsMap, name);
+        }
         PythonLikeObject out = globalsMap.get(name);
         if (out == null) {
             return GlobalBuiltins.lookupOrError(this, name);

jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/util/PythonGlobalsBackedMap.java

@@ -0,0 +1,17 @@
+package ai.timefold.jpyinterpreter.util;
+
+import java.util.HashMap;
+
+import ai.timefold.jpyinterpreter.PythonLikeObject;
+
+public class PythonGlobalsBackedMap extends HashMap<String, PythonLikeObject> {
+    private final long pythonGlobalsId;
+
+    public PythonGlobalsBackedMap(long pythonGlobalsId) {
+        this.pythonGlobalsId = pythonGlobalsId;
+    }
+
+    public long getPythonGlobalsId() {
+        return pythonGlobalsId;
+    }
+}

jpyinterpreter/src/main/python/jvm_setup.py

@@ -41,6 +41,7 @@ def init(*args, path: List[str] = None, include_translator_jars: bool = True,
     CPythonBackedPythonInterpreter.lookupPythonReferenceIdPythonFunction = GetPythonObjectId()
     CPythonBackedPythonInterpreter.lookupPythonReferenceTypePythonFunction = GetPythonObjectType()
     CPythonBackedPythonInterpreter.lookupAttributeOnPythonReferencePythonFunction = GetAttributeOnPythonObject()
+    CPythonBackedPythonInterpreter.loadObjectFromPythonGlobalDict = GetNameFromGlobals()
     CPythonBackedPythonInterpreter.lookupPointerForAttributeOnPythonReferencePythonFunction = \
         GetAttributePointerOnPythonObject()
     CPythonBackedPythonInterpreter.lookupPointerArrayForAttributeOnPythonReferencePythonFunction = \
@@ -55,6 +56,25 @@ def init(*args, path: List[str] = None, include_translator_jars: bool = True,
     CPythonBackedPythonInterpreter.importModuleFunction = ImportModule()
 
 
+@jpype.JImplements('java.util.function.BiConsumer', deferred=True)
+class GetNameFromGlobals:
+    @jpype.JOverride()
+    def accept(self, java_globals, name):
+        from .translator import java_globals_to_python_globals
+        from .conversions import convert_to_java_python_like_object
+        from ai.timefold.jpyinterpreter.util import PythonGlobalsBackedMap
+
+        if not isinstance(java_globals, PythonGlobalsBackedMap):
+            return
+
+        python_globals = java_globals_to_python_globals[java_globals.getPythonGlobalsId()]
+        try:
+            python_object = python_globals[name]
+            java_globals.put(name, convert_to_java_python_like_object(python_object))
+        except KeyError:
+            java_globals.put(name, None)
+
+
 @jpype.JImplements('java.util.function.Function', deferred=True)
 class GetPythonObjectId:
     @jpype.JOverride()

jpyinterpreter/src/main/python/translator.py

@@ -13,6 +13,8 @@
 
 global_dict_to_instance = dict()
 global_dict_to_key_set = dict()
+java_globals_to_python_globals = dict()
+
 type_to_compiled_java_class = dict()
 type_to_annotations = dict()
 type_to_java_interfaces = dict()
@@ -136,26 +138,33 @@ def copy_closure(closure):
         return out
 
 
-def copy_globals(globals_dict, co_names):
+def copy_globals(globals_dict, co_names, python_class):
     global global_dict_to_instance
     global global_dict_to_key_set
     from .conversions import convert_to_java_python_like_object
-    from java.util import HashMap
+    from ai.timefold.jpyinterpreter.util import PythonGlobalsBackedMap
     from ai.timefold.jpyinterpreter import CPythonBackedPythonInterpreter
 
     globals_dict_key = id(globals_dict)
     if globals_dict_key in global_dict_to_instance:
         out = global_dict_to_instance[globals_dict_key]
         key_set = global_dict_to_key_set[globals_dict_key]
     else:
-        out = HashMap()
+        out = PythonGlobalsBackedMap(globals_dict_key)
         key_set = set()
         global_dict_to_instance[globals_dict_key] = out
         global_dict_to_key_set[globals_dict_key] = key_set
+        java_globals_to_python_globals[globals_dict_key] = globals_dict
 
     instance_map = CPythonBackedPythonInterpreter.pythonObjectIdToConvertedObjectMap
     for key, value in globals_dict.items():
         if key not in key_set and key in co_names:
+            if python_class is not None:
+                if isinstance(value, type):
+                    if issubclass(value, python_class):
+                        continue
+                elif isinstance(value, python_class):
+                    continue
             key_set.add(key)
             out.put(key, convert_to_java_python_like_object(value, instance_map))
     return out
@@ -216,7 +225,7 @@ def get_python_exception_table(python_code):
     return out
 
 
-def get_function_bytecode_object(python_function):
+def get_function_bytecode_object(python_function, python_class: type = None):
     from .annotations import copy_type_annotations
     from .conversions import copy_iterable, init_type_to_compiled_java_class, convert_to_java_python_like_object
     from java.util import ArrayList
@@ -251,7 +260,8 @@ def get_function_bytecode_object(python_function):
     python_compiled_function.co_argcount = python_function.__code__.co_argcount
     python_compiled_function.co_kwonlyargcount = python_function.__code__.co_kwonlyargcount
     python_compiled_function.closure = copy_closure(python_function.__closure__)
-    python_compiled_function.globalsMap = copy_globals(python_function.__globals__, python_function.__code__.co_names)
+    python_compiled_function.globalsMap = copy_globals(python_function.__globals__, python_function.__code__.co_names,
+                                                       python_class)
     python_compiled_function.typeAnnotations = copy_type_annotations(python_function,
                                                                      get_default_args(python_function),
                                                                      inspect.getfullargspec(python_function).varargs,
@@ -267,7 +277,7 @@ def get_function_bytecode_object(python_function):
 
 
 def get_static_function_bytecode_object(the_class, python_function):
-    return get_function_bytecode_object(python_function.__get__(the_class))
+    return get_function_bytecode_object(python_function.__get__(the_class), python_class=the_class)
 
 
 def copy_variable_names(iterable):
@@ -673,7 +683,7 @@ def translate_python_class_to_java_class(python_class):
 
     instance_method_map = HashMap()
     for method in instance_methods:
-        instance_method_map.put(method[0], get_function_bytecode_object(method[1]))
+        instance_method_map.put(method[0], get_function_bytecode_object(method[1], python_class=python_class))
 
     static_attributes_map = HashMap()
     static_attributes_to_class_instance_map = HashMap()
@@ -683,7 +693,7 @@ def translate_python_class_to_java_class(python_class):
 
     for attribute in static_attributes:
         attribute_type = type(attribute[1])
-        if attribute_type == python_class:
+        if issubclass(attribute_type, python_class):
             static_attributes_to_class_instance_map.put(attribute[0],
                                                         JProxy(OpaquePythonReference,
                                                                inst=attribute[1], convert=True))

jpyinterpreter/tests/test_classes.py

@@ -896,6 +896,20 @@ def my_function_instance(a, b, c):
     verifier.verify(1, 1, 1, expected_result=13)
 
 
+def test_enum_translate_to_class():
+    from enum import Enum
+    from jpyinterpreter import translate_python_class_to_java_class
+    from ai.timefold.jpyinterpreter.types.wrappers import CPythonType
+
+    class Color(Enum):
+        RED = 'RED'
+        GREEN = 'GREEN'
+        BLUE = 'BLUE'
+
+    translated_class = translate_python_class_to_java_class(Color)
+    assert not isinstance(translated_class, CPythonType)
+
+
 def test_class_annotations():
     from typing import Annotated
     from java.lang import Deprecated