routing(python): fix SetAllowedVehiclesForIndex (#4982)

Author: Mizux

ortools/constraint_solver/python/pywraprouting_test.py

@@ -564,6 +564,33 @@ def testMatrixDimensionVRP(self):
                     manager.IndexToNode(index)
                 ]
 
+    def testAllowedVehicles(self):
+        manager = pywrapcp.RoutingIndexManager(
+            2, 2, 0  # num_nodes  # num_vehicles  # depot
+        )
+        self.assertIsNotNone(manager)
+        model = pywrapcp.RoutingModel(manager)
+        self.assertIsNotNone(model)
+        # out of range vehicle index is allowed.
+        model.SetAllowedVehiclesForIndex(vehicles=[13], index=1)
+        self.assertFalse(model.IsVehicleAllowedForIndex(vehicle=0, index=1))
+        self.assertFalse(model.IsVehicleAllowedForIndex(1, 1))
+        self.assertTrue(model.IsVehicleAllowedForIndex(13, 1))
+        # empty list means any vehicles are allowed.
+        model.SetAllowedVehiclesForIndex([], 1)
+        self.assertTrue(model.IsVehicleAllowedForIndex(0, 1))
+        self.assertTrue(model.IsVehicleAllowedForIndex(1, 1))
+        self.assertTrue(model.IsVehicleAllowedForIndex(42, 1))
+        # only allow first vehicle for node 1.
+        model.SetAllowedVehiclesForIndex([0], 1)
+        self.assertTrue(model.IsVehicleAllowedForIndex(0, 1))
+        self.assertFalse(model.IsVehicleAllowedForIndex(1, 1))
+        self.assertFalse(model.IsVehicleAllowedForIndex(2, 1))
+        assignment = model.Solve()
+        self.assertIsNotNone(assignment)
+        self.assertEqual(1, assignment.Value(model.NextVar(model.Start(0))))
+        self.assertEqual(model.GetVehicleClassesCount(), 2)
+
     def testDisjunctionTSP(self):
         # Create routing model
         manager = pywrapcp.RoutingIndexManager(10, 1, 0)

ortools/constraint_solver/python/routing.swig

@@ -71,6 +71,14 @@ DEFINE_INDEX_TYPE_TYPEDEF(
 %ignore operations_research::RoutingModel::AddDimensionDependentDimensionWithVehicleCapacity;
 %ignore operations_research::RoutingModel::AddResourceGroup;
 %ignore operations_research::RoutingModel::GetResourceGroups;
+%ignore operations_research::RoutingModel::SetAllowedVehiclesForIndex(absl::Span<const int>, int64_t);
+%extend operations_research::RoutingModel {
+  void SetAllowedVehiclesForIndex(
+    const std::vector<int>& vehicles,
+    int64_t index) {
+      $self->SetAllowedVehiclesForIndex(absl::Span<const int>(vehicles), index);
+    }
+}
 
 PY_PROTO_TYPEMAP(ortools.constraint_solver.routing_parameters_pb2,
                  RoutingModelParameters,