gcode_move: Support additional toolhead axes

Signed-off-by: Kevin O'Connor <kevin@koconnor.net>

1 files changed, 22 insertions, 6 deletions

diff --git a/klippy/extras/gcode_move.py b/klippy/extras/gcode_move.py
index d5930dc2..94a0ce42 100644
--- a/klippy/extras/gcode_move.py
+++ b/klippy/extras/gcode_move.py
@@ -14,6 +14,8 @@ class GCodeMove:
                                        self.reset_last_position)
         printer.register_event_handler("toolhead:manual_move",
                                        self.reset_last_position)
+        printer.register_event_handler("toolhead:update_extra_axes",
+                                       self._update_extra_axes)
         printer.register_event_handler("gcode:command_error",
                                        self.reset_last_position)
         printer.register_event_handler("extruder:activate_extruder",
@@ -103,13 +105,27 @@ class GCodeMove:
             'extrude_factor': self.extrude_factor,
             'absolute_coordinates': self.absolute_coord,
             'absolute_extrude': self.absolute_extrude,
-            'homing_origin': self.Coord(*self.homing_position),
-            'position': self.Coord(*self.last_position),
-            'gcode_position': self.Coord(*move_position),
+            'homing_origin': self.Coord(*self.homing_position[:4]),
+            'position': self.Coord(*self.last_position[:4]),
+            'gcode_position': self.Coord(*move_position[:4]),
         }
     def reset_last_position(self):
         if self.is_printer_ready:
             self.last_position = self.position_with_transform()
+    def _update_extra_axes(self):
+        toolhead = self.printer.lookup_object('toolhead')
+        axis_map = {'X':0, 'Y': 1, 'Z': 2, 'E': 3}
+        extra_axes = toolhead.get_extra_axes()
+        for index, ea in enumerate(extra_axes):
+            if ea is None:
+                continue
+            gcode_id = ea.get_axis_gcode_id()
+            if gcode_id is None or gcode_id in axis_map or gcode_id in "FN":
+                continue
+            axis_map[gcode_id] = index
+        self.axis_map = axis_map
+        self.base_position[4:] = [0.] * (len(extra_axes) - 4)
+        self.reset_last_position()
     # G-Code movement commands
     def cmd_G1(self, gcmd):
         # Move
@@ -167,7 +183,7 @@ class GCodeMove:
                     offset *= self.extrude_factor
                 self.base_position[i] = self.last_position[i] - offset
         if offsets == [None, None, None, None]:
-            self.base_position = list(self.last_position)
+            self.base_position[:4] = self.last_position[:4]
     def cmd_M114(self, gcmd):
         # Get Current Position
         p = self._get_gcode_position()
@@ -225,7 +241,7 @@ class GCodeMove:
         # Restore state
         self.absolute_coord = state['absolute_coord']
         self.absolute_extrude = state['absolute_extrude']
-        self.base_position = list(state['base_position'])
+        self.base_position[:4] = state['base_position'][:4]
         self.homing_position = list(state['homing_position'])
         self.speed = state['speed']
         self.speed_factor = state['speed_factor']
@@ -253,7 +269,7 @@ class GCodeMove:
         kinfo = zip("XYZ", kin.calc_position(dict(cinfo)))
         kin_pos = " ".join(["%s:%.6f" % (a, v) for a, v in kinfo])
         toolhead_pos = " ".join(["%s:%.6f" % (a, v) for a, v in zip(
-            "XYZE", toolhead.get_position())])
+            "XYZE", toolhead.get_position()[:4])])
         gcode_pos = " ".join(["%s:%.6f"  % (a, v)
                               for a, v in zip("XYZE", self.last_position)])
         base_pos = " ".join(["%s:%.6f"  % (a, v)