feat: Add support for nearby selection (#32)

Christopher-Chianelli · web-flow · commit 30a13b982291 · 2024-04-16T10:26:30.000-04:00
- Add automatic type coercion between int and float to prevent
  ClassCastExceptions when a function expecting float is given
  int and vice-versa
- Add register_java_class to set field holding the generated class
  and register the class in the SolverConfig classloader
- Rename field holding the generated class from __timefold_java_class
  to _timefold_java_class to prevent name-mangling
- Remove unused code
diff --git a/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/PythonClassTranslator.java b/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/PythonClassTranslator.java
@@ -23,6 +23,7 @@
 import ai.timefold.jpyinterpreter.implementors.JavaEqualsImplementor;
 import ai.timefold.jpyinterpreter.implementors.JavaHashCodeImplementor;
 import ai.timefold.jpyinterpreter.implementors.JavaInterfaceImplementor;
+import ai.timefold.jpyinterpreter.implementors.JavaPythonTypeConversionImplementor;
 import ai.timefold.jpyinterpreter.implementors.PythonConstantsImplementor;
 import ai.timefold.jpyinterpreter.opcodes.AbstractOpcode;
 import ai.timefold.jpyinterpreter.opcodes.Opcode;
@@ -1096,6 +1097,13 @@ public static void createSetAttribute(ClassWriter classWriter, String classInter
                 getUnwrappedJavaObject(methodVisitor, type);
                 typeDescriptor = Type.getDescriptor(type.getJavaObjectWrapperType());
             } else {
+                methodVisitor.visitLdcInsn(Type.getType(type.getJavaTypeDescriptor()));
+                methodVisitor.visitMethodInsn(Opcodes.INVOKESTATIC,
+                        Type.getInternalName(JavaPythonTypeConversionImplementor.class),
+                        "coerceToType", Type.getMethodDescriptor(Type.getType(Object.class),
+                                Type.getType(PythonLikeObject.class),
+                                Type.getType(Class.class)),
+                        false);
                 methodVisitor.visitTypeInsn(Opcodes.CHECKCAST, type.getJavaTypeInternalName());
             }
             methodVisitor.visitFieldInsn(Opcodes.PUTFIELD, classInternalName, getJavaFieldName(field),
@@ -1231,6 +1239,12 @@ public static void createReadFromCPythonReference(ClassWriter classWriter, Strin
             }
 
             var attributeType = attributeNameToType.get(field);
+            methodVisitor.visitLdcInsn(Type.getType(attributeType.getJavaTypeDescriptor()));
+            methodVisitor.visitMethodInsn(Opcodes.INVOKESTATIC, Type.getInternalName(JavaPythonTypeConversionImplementor.class),
+                    "coerceToType", Type.getMethodDescriptor(Type.getType(Object.class),
+                            Type.getType(PythonLikeObject.class),
+                            Type.getType(Class.class)),
+                    false);
             methodVisitor.visitTypeInsn(Opcodes.CHECKCAST, attributeNameToType.get(field).getJavaTypeInternalName());
 
             if (attributeType.getJavaTypeInternalName().equals(Type.getInternalName(JavaObjectWrapper.class))) {
diff --git a/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/implementors/JavaPythonTypeConversionImplementor.java b/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/implementors/JavaPythonTypeConversionImplementor.java
@@ -14,6 +14,7 @@
 import ai.timefold.jpyinterpreter.PythonLikeObject;
 import ai.timefold.jpyinterpreter.StackMetadata;
 import ai.timefold.jpyinterpreter.types.BuiltinTypes;
+import ai.timefold.jpyinterpreter.types.Coercible;
 import ai.timefold.jpyinterpreter.types.PythonByteArray;
 import ai.timefold.jpyinterpreter.types.PythonBytes;
 import ai.timefold.jpyinterpreter.types.PythonCode;
@@ -447,6 +448,31 @@ public static void returnValue(MethodVisitor methodVisitor, MethodDescriptor met
         methodVisitor.visitInsn(Opcodes.ARETURN);
     }
 
+    /**
+     * Coerce a value to a given type
+     */
+    public static <T> T coerceToType(PythonLikeObject value, Class<T> type) {
+        if (value == null) {
+            return null;
+        }
+
+        if (type.isAssignableFrom(value.getClass())) {
+            return (T) value;
+        }
+
+        if (value instanceof Coercible coercible) {
+            var out = coercible.coerce(type);
+            if (out == null) {
+                throw new TypeError("%s cannot be coerced to %s."
+                        .formatted(value.$getType(), type));
+            }
+            return out;
+        }
+
+        throw new TypeError("%s cannot be coerced to %s."
+                .formatted(value.$getType(), type));
+    }
+
     /**
      * Convert the {@code parameterIndex} Java parameter to its Python equivalent and store it into
      * the corresponding Python parameter local variable slot.
diff --git a/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/implementors/KnownCallImplementor.java b/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/implementors/KnownCallImplementor.java
@@ -113,6 +113,12 @@ public static void callMethod(PythonFunctionSignature pythonFunctionSignature, M
         // Now load and typecheck the local variables
         for (int i = 0; i < Math.min(specPositionalArgumentCount, argumentCount); i++) {
             localVariableHelper.readTemp(methodVisitor, Type.getType(PythonLikeObject.class), argumentLocals[i]);
+            methodVisitor.visitLdcInsn(Type.getType(pythonFunctionSignature.getArgumentSpec().getArgumentType(i)));
+            methodVisitor.visitMethodInsn(Opcodes.INVOKESTATIC, Type.getInternalName(JavaPythonTypeConversionImplementor.class),
+                    "coerceToType", Type.getMethodDescriptor(Type.getType(Object.class),
+                            Type.getType(PythonLikeObject.class),
+                            Type.getType(Class.class)),
+                    false);
             methodVisitor.visitTypeInsn(Opcodes.CHECKCAST,
                     Type.getInternalName(pythonFunctionSignature.getArgumentSpec().getArgumentType(i)));
         }
@@ -279,6 +285,12 @@ public static void callPython311andAbove(PythonFunctionSignature pythonFunctionS
         // Load arguments in proper order and typecast them
         for (int i = 0; i < specTotalArgumentCount; i++) {
             localVariableHelper.readTemp(methodVisitor, Type.getType(PythonLikeObject.class), argumentLocals[i]);
+            methodVisitor.visitLdcInsn(Type.getType(pythonFunctionSignature.getArgumentSpec().getArgumentType(i)));
+            methodVisitor.visitMethodInsn(Opcodes.INVOKESTATIC, Type.getInternalName(JavaPythonTypeConversionImplementor.class),
+                    "coerceToType", Type.getMethodDescriptor(Type.getType(Object.class),
+                            Type.getType(PythonLikeObject.class),
+                            Type.getType(Class.class)),
+                    false);
             methodVisitor.visitTypeInsn(Opcodes.CHECKCAST,
                     Type.getInternalName(pythonFunctionSignature.getArgumentSpec().getArgumentType(i)));
         }
@@ -480,6 +492,12 @@ public static void callUnpackListAndMap(Class<?> defaultArgumentHolderClass, Met
             methodVisitor.visitMethodInsn(Opcodes.INVOKEINTERFACE, Type.getInternalName(List.class),
                     "get", Type.getMethodDescriptor(Type.getType(Object.class), Type.INT_TYPE),
                     true);
+            methodVisitor.visitLdcInsn(descriptorParameterTypes[i]);
+            methodVisitor.visitMethodInsn(Opcodes.INVOKESTATIC, Type.getInternalName(JavaPythonTypeConversionImplementor.class),
+                    "coerceToType", Type.getMethodDescriptor(Type.getType(Object.class),
+                            Type.getType(PythonLikeObject.class),
+                            Type.getType(Class.class)),
+                    false);
             methodVisitor.visitTypeInsn(Opcodes.CHECKCAST, descriptorParameterTypes[i].getInternalName());
             methodVisitor.visitInsn(Opcodes.SWAP);
         }
diff --git a/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/implementors/ObjectImplementor.java b/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/implementors/ObjectImplementor.java
@@ -140,6 +140,12 @@ public static void setAttribute(FunctionMetadata functionMetadata, MethodVisitor
             FieldDescriptor fieldDescriptor = maybeFieldDescriptor.get();
             methodVisitor.visitTypeInsn(Opcodes.CHECKCAST, fieldDescriptor.declaringClassInternalName());
             StackManipulationImplementor.swap(methodVisitor);
+            methodVisitor.visitLdcInsn(Type.getType(fieldDescriptor.fieldPythonLikeType().getJavaTypeDescriptor()));
+            methodVisitor.visitMethodInsn(Opcodes.INVOKESTATIC, Type.getInternalName(JavaPythonTypeConversionImplementor.class),
+                    "coerceToType", Type.getMethodDescriptor(Type.getType(Object.class),
+                            Type.getType(PythonLikeObject.class),
+                            Type.getType(Class.class)),
+                    false);
             methodVisitor.visitTypeInsn(Opcodes.CHECKCAST, fieldDescriptor.fieldPythonLikeType().getJavaTypeInternalName());
             if (fieldDescriptor.isJavaType()) {
                 // Need to unwrap the object
diff --git a/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/types/Coercible.java b/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/types/Coercible.java
@@ -0,0 +1,5 @@
+package ai.timefold.jpyinterpreter.types;
+
+public interface Coercible {
+    <T> T coerce(Class<T> targetType);
+}
diff --git a/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/types/numeric/PythonFloat.java b/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/types/numeric/PythonFloat.java
@@ -15,6 +15,7 @@
 import ai.timefold.jpyinterpreter.PythonUnaryOperator;
 import ai.timefold.jpyinterpreter.types.AbstractPythonLikeObject;
 import ai.timefold.jpyinterpreter.types.BuiltinTypes;
+import ai.timefold.jpyinterpreter.types.Coercible;
 import ai.timefold.jpyinterpreter.types.NotImplemented;
 import ai.timefold.jpyinterpreter.types.PythonLikeComparable;
 import ai.timefold.jpyinterpreter.types.PythonLikeFunction;
@@ -29,7 +30,8 @@
 import ai.timefold.jpyinterpreter.util.StringFormatter;
 import ai.timefold.solver.core.impl.domain.solution.cloner.PlanningImmutable;
 
-public class PythonFloat extends AbstractPythonLikeObject implements PythonNumber, PlanningImmutable {
+public class PythonFloat extends AbstractPythonLikeObject implements PythonNumber, PlanningImmutable,
+        Coercible {
     public final double value;
 
     static {
@@ -793,4 +795,12 @@ private DecimalFormat getNumberFormat(DefaultFormatSpec formatSpec) {
         StringFormatter.align(out, formatSpec, DefaultFormatSpec.AlignmentOption.RIGHT_ALIGN);
         return PythonString.valueOf(out.toString());
     }
+
+    @Override
+    public <T> T coerce(Class<T> targetType) {
+        if (targetType.equals(PythonInteger.class)) {
+            return (T) PythonInteger.valueOf((long) value);
+        }
+        return null;
+    }
 }
diff --git a/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/types/numeric/PythonInteger.java b/jpyinterpreter/src/main/java/ai/timefold/jpyinterpreter/types/numeric/PythonInteger.java
@@ -13,6 +13,7 @@
 import ai.timefold.jpyinterpreter.PythonUnaryOperator;
 import ai.timefold.jpyinterpreter.types.AbstractPythonLikeObject;
 import ai.timefold.jpyinterpreter.types.BuiltinTypes;
+import ai.timefold.jpyinterpreter.types.Coercible;
 import ai.timefold.jpyinterpreter.types.NotImplemented;
 import ai.timefold.jpyinterpreter.types.PythonLikeFunction;
 import ai.timefold.jpyinterpreter.types.PythonLikeType;
@@ -26,7 +27,8 @@
 import ai.timefold.jpyinterpreter.util.StringFormatter;
 import ai.timefold.solver.core.impl.domain.solution.cloner.PlanningImmutable;
 
-public class PythonInteger extends AbstractPythonLikeObject implements PythonNumber, PlanningImmutable {
+public class PythonInteger extends AbstractPythonLikeObject implements PythonNumber,
+        PlanningImmutable, Coercible {
     private static final BigInteger MIN_BYTE = BigInteger.valueOf(0);
     private static final BigInteger MAX_BYTE = BigInteger.valueOf(255);
 
@@ -840,4 +842,12 @@ public PythonString asString() {
 
         return PythonString.valueOf(out.toString());
     }
+
+    @Override
+    public <T> T coerce(Class<T> targetType) {
+        if (targetType.equals(PythonFloat.class)) {
+            return (T) PythonFloat.valueOf(value.doubleValue());
+        }
+        return null;
+    }
 }
diff --git a/jpyinterpreter/tests/test_classes.py b/jpyinterpreter/tests/test_classes.py
@@ -24,6 +24,28 @@ def create_instance(x: int) -> A:
     verifier.verify(3, expected_result=A(3))
 
 
+def test_type_coercing():
+    class A:
+        value: float
+
+        def __init__(self, value):
+            self.value = value
+
+        def __eq__(self, other):
+            if not isinstance(other, A):
+                return False
+            return self.value == other.value
+
+    def create_instance(x: int) -> A:
+        return A(x)
+
+    verifier = verifier_for(create_instance)
+
+    verifier.verify(1, expected_result=A(1))
+    verifier.verify(2, expected_result=A(2))
+    verifier.verify(3, expected_result=A(3))
+
+
 def test_deleted_field():
     class A:
         value: int
diff --git a/tests/test_solver_config.py b/tests/test_solver_config.py
@@ -55,7 +55,7 @@ def test_load_from_solver_config_file():
     assert entity_class_list.size() == 1
     assert entity_class_list.get(0) == get_java_type_for_python_type(Entity).getJavaClass()
     assert solver_config.getScoreDirectorFactoryConfig().getConstraintProviderClass() == \
-           my_constraints.__timefold_java_class  # noqa
+           my_constraints._timefold_java_class  # noqa
     assert solver_config.getTerminationConfig().getBestScoreLimit() == '0hard/0soft'
 
 
diff --git a/tests/test_user_error.py b/tests/test_user_error.py
@@ -111,3 +111,12 @@ def test_missing_enterprise():
         solver_config = SolverConfig(
             move_thread_count=MoveThreadCount.AUTO
         )._to_java_solver_config()
+
+    @nearby_distance_meter
+    def my_distance_meter(entity: Entity, value: str) -> float:
+        return 0.0
+
+    with pytest.raises(RequiresEnterpriseError, match=re.escape('nearby selection')):
+        solver_config = SolverConfig(
+            nearby_distance_meter_function=my_distance_meter
+        )._to_java_solver_config()
diff --git a/timefold-solver-python-core/src/main/python/annotation/annotations.py b/timefold-solver-python-core/src/main/python/annotation/annotations.py
@@ -1,23 +1,23 @@
 import jpype
 
-from ..timefold_java_interop import ensure_init, _generate_constraint_provider_class, \
-    _generate_incremental_score_calculator_class, \
-    _generate_variable_listener_class
+from ..constraint import ConstraintFactory
+from ..timefold_java_interop import ensure_init, _generate_constraint_provider_class, register_java_class
 from jpyinterpreter import JavaAnnotation
 from jpype import JImplements, JOverride
 from typing import Union, List, Callable, Type, TYPE_CHECKING, TypeVar
 
 if TYPE_CHECKING:
     from ai.timefold.solver.core.api.solver.change import ProblemChange as _ProblemChange
-    from ai.timefold.solver.core.api.score.stream import (Constraint as _Constraint,
-                                                          ConstraintFactory as _ConstraintFactory)
+    from ai.timefold.solver.core.api.score.stream import Constraint as _Constraint
     from ai.timefold.solver.core.api.score import Score as _Score
     from ai.timefold.solver.core.api.score.calculator import IncrementalScoreCalculator as _IncrementalScoreCalculator
     from ai.timefold.solver.core.api.domain.variable import PlanningVariableGraphType as _PlanningVariableGraphType, \
         VariableListener as _VariableListener
 
 
 Solution_ = TypeVar('Solution_')
+Origin_ = TypeVar('Origin_')
+Destination_ = TypeVar('Destination_')
 
 
 class PlanningId(JavaAnnotation):
@@ -301,10 +301,20 @@ def __init__(self, a_list):
     return out
 
 
-def constraint_provider(constraint_provider_function: Callable[['_ConstraintFactory'], List['_Constraint']] = None, /) \
-        -> \
-        Union[Callable[[Callable], Callable[['_ConstraintFactory'], List['_Constraint']]],
-              Callable[['_ConstraintFactory'], List['_Constraint']]]:
+def nearby_distance_meter(distance_function: Callable[[Origin_, Destination_], float], /) \
+        -> Callable[[Origin_, Destination_], float]:
+    ensure_init()
+    from jpyinterpreter import translate_python_bytecode_to_java_bytecode, generate_proxy_class_for_translated_function
+    from ai.timefold.solver.core.impl.heuristic.selector.common.nearby import NearbyDistanceMeter  # noqa
+    java_class = generate_proxy_class_for_translated_function(NearbyDistanceMeter,
+                                                              translate_python_bytecode_to_java_bytecode(
+                                                                  distance_function,
+                                                                  NearbyDistanceMeter))
+    return register_java_class(distance_function, java_class)
+
+
+def constraint_provider(constraint_provider_function: Callable[[ConstraintFactory], List['_Constraint']], /) \
+        -> Callable[[ConstraintFactory], List['_Constraint']]:
     """Marks a function as a ConstraintProvider.
 
     The function takes a single parameter, the ConstraintFactory, and
@@ -321,9 +331,8 @@ def wrapped_constraint_provider(constraint_factory):
             from ..constraint import ConstraintFactory
             out = function(ConstraintFactory(constraint_factory))
             return out
-        wrapped_constraint_provider.__timefold_java_class = (
-            _generate_constraint_provider_class(function, wrapped_constraint_provider))
-        return wrapped_constraint_provider
+        java_class = _generate_constraint_provider_class(function, wrapped_constraint_provider)
+        return register_java_class(wrapped_constraint_provider, java_class)
 
     return constraint_provider_wrapper(constraint_provider_function)
 
@@ -342,10 +351,10 @@ def easy_score_calculator(easy_score_calculator_function: Callable[[Solution_],
     ensure_init()
     from jpyinterpreter import translate_python_bytecode_to_java_bytecode, generate_proxy_class_for_translated_function
     from ai.timefold.solver.core.api.score.calculator import EasyScoreCalculator
-    easy_score_calculator_function.__timefold_java_class = \
-        generate_proxy_class_for_translated_function(EasyScoreCalculator,
-            translate_python_bytecode_to_java_bytecode(easy_score_calculator_function, EasyScoreCalculator))
-    return easy_score_calculator_function
+    java_class = generate_proxy_class_for_translated_function(EasyScoreCalculator,
+                                                              translate_python_bytecode_to_java_bytecode(
+                                                                  easy_score_calculator_function, EasyScoreCalculator))
+    return register_java_class(easy_score_calculator_function, java_class)
 
 
 def incremental_score_calculator(incremental_score_calculator: Type['_IncrementalScoreCalculator']) -> \
@@ -415,9 +424,8 @@ def resetWorkingSolution(self, workingSolution: Solution_, constraintMatchEnable
     #     method_on_class = getattr(incremental_score_calculator, method, None
 
     translated_class = translate_python_class_to_java_class(incremental_score_calculator)
-    incremental_score_calculator.__timefold_java_class = generate_proxy_class_for_translated_class(base_interface,
-                                                                                                   translated_class)
-    return incremental_score_calculator
+    java_class = generate_proxy_class_for_translated_class(base_interface, translated_class)
+    return register_java_class(incremental_score_calculator, java_class)
 
 
 def variable_listener(variable_listener_class: Type['_VariableListener'] = None, /, *,
@@ -505,9 +513,8 @@ def reset_working_solution(self, score_director):
             setattr(the_variable_listener_class, 'resetWorkingSolution', reset_working_solution)
 
         translated_class = translate_python_class_to_java_class(the_variable_listener_class)
-        the_variable_listener_class.__timefold_java_class = generate_proxy_class_for_translated_class(VariableListener,
-                                                                                                      translated_class)
-        return the_variable_listener_class
+        java_class = generate_proxy_class_for_translated_class(VariableListener, translated_class)
+        return register_java_class(the_variable_listener_class, java_class)
 
     if variable_listener_class:  # Called as @variable_listener
         return variable_listener_wrapper(variable_listener_class)
@@ -570,5 +577,6 @@ def wrapper_doChange(self, solution, problem_change_director):
            'PlanningEntityProperty', 'PlanningEntityCollectionProperty',
            'ValueRangeProvider', 'DeepPlanningClone',
            'planning_entity', 'planning_solution',
+           'nearby_distance_meter',
            'constraint_provider', 'easy_score_calculator', 'incremental_score_calculator',
            'variable_listener', 'problem_change']
diff --git a/timefold-solver-python-core/src/main/python/config/config.py b/timefold-solver-python-core/src/main/python/config/config.py
diff --git a/timefold-solver-python-core/src/main/python/jpype_type_conversions.py b/timefold-solver-python-core/src/main/python/jpype_type_conversions.py
diff --git a/timefold-solver-python-core/src/main/python/timefold_java_interop.py b/timefold-solver-python-core/src/main/python/timefold_java_interop.py
