From 7431231c0cbda64e10ef47e56f11f87d1afbcfa4 Mon Sep 17 00:00:00 2001
From: Jonathan Feinberg <jonathf@gmail.com>
Date: Fri, 19 Apr 2024 15:25:23 +0200
Subject: [PATCH] trying it out

---
 src/bind/solution/solution.cpp |  7 +++++++
 src/vroom/input/input.py       | 10 +++++++++-
 src/vroom/solution/solution.py | 14 ++++++++++++--
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/bind/solution/solution.cpp b/src/bind/solution/solution.cpp
index da0d1b9..d392f1f 100644
--- a/src/bind/solution/solution.cpp
+++ b/src/bind/solution/solution.cpp
@@ -90,6 +90,7 @@ void init_solution(py::module_ &m) {
                  ptr[idx].setup = step.setup;
                  ptr[idx].service = step.service;
                  ptr[idx].waiting_time = step.waiting_time;
+                 ptr[idx].distance = step.distance;
                  ptr[idx].arrival = step.arrival;
                  ptr[idx].duration = step.duration;
 
@@ -99,6 +100,12 @@ void init_solution(py::module_ &m) {
              return arr;
            })
       .def("_solution_json",
+           [](vroom::Solution solution) {
+             py::scoped_ostream_redirect stream(
+                 std::cout, py::module_::import("sys").attr("stdout"));
+             vroom::io::write_to_json(solution, false, "");
+           })
+      .def("_geometry_solution_json",
            [](vroom::Solution solution) {
              py::scoped_ostream_redirect stream(
                  std::cout, py::module_::import("sys").attr("stdout"));
diff --git a/src/vroom/input/input.py b/src/vroom/input/input.py
index bbf3f2e..d989ad1 100644
--- a/src/vroom/input/input.py
+++ b/src/vroom/input/input.py
@@ -29,6 +29,8 @@ class Input(_vroom.Input):
 
     """
 
+    _geometry: bool = False
+
     def __init__(
         self,
         amount_size: Optional[int] = None,
@@ -104,12 +106,15 @@ def from_json(
         """
         if geometry is None:
             geometry = servers is not None
+        if geometry:
+            self._set_geometry()
         instance = Input(servers=servers, router=router)
         with open(filepath) as handle:
             instance._from_json(handle.read(), geometry)
         return instance
 
     def set_geometry(self):
+        self._geometry = True
         return self._set_geometry()
 
     def set_amount_size(self, *amount_sizes: int) -> None:
@@ -307,9 +312,12 @@ def solve(
         exploration_level: int,
         nb_threads: int,
     ) -> Solution:
-        return Solution(
+        solution = Solution(
             self._solve(
                 exploration_level=exploration_level,
                 nb_threads=nb_threads,
             )
+
         )
+        solution._geometry = self._geometry
+        return solution
diff --git a/src/vroom/solution/solution.py b/src/vroom/solution/solution.py
index 7faea0b..7a2c0ba 100644
--- a/src/vroom/solution/solution.py
+++ b/src/vroom/solution/solution.py
@@ -22,6 +22,8 @@ class Solution(_vroom.Solution):
             Frame outlining all routes for all vehicles.
     """
 
+    _geometry: bool = False
+
     @property
     def routes(self) -> pandas.DataFrame:
         """
@@ -81,17 +83,25 @@ def routes(self) -> pandas.DataFrame:
                 del frame[column]
             else:
                 frame.loc[frame[column] == NA_SUBSTITUTE, column] = pandas.NA
+        if self._geometry:
+            frame["distance"] = array["distance"]
         return frame
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert solution into VROOM compatible dictionary."""
         stream = io.StringIO()
         with redirect_stdout(stream):
-            self._solution_json()
+            if self._geometry:
+                self._geometry_solution_json()
+            else:
+                self._solution_json()
         return json.loads(stream.getvalue())
 
     def to_json(self, filepath: Union[str, Path]) -> None:
         """Store solution into VROOM compatible JSON file."""
         with open(filepath, "w") as handler:
             with redirect_stdout(handler):
-                self._solution_json()
+                if self._geometry:
+                    self._geometry_solution_json()
+                else:
+                    self._solution_json()