Raw Text Content QR
falcon-spider-pig



# Interactively identify printer axis motors via web ui prompts
#
# Copyright (C) 2026  the Klipper contributors
#
# This file may be distributed under the terms of the GNU GPLv3 license.
import glob, logging, os, re, time

AXES = "xyz"

STEPPER_NAME_R = re.compile(r'(?i)^stepper_([a-z])(\d*)$')
# Sections that follow a stepper by name (eg, "[tmc2209 stepper_x]")
DEP_SECTION_R = re.compile(r'(?i)^(tmc\d+|endstop_phase)\s+(stepper_\S+)$')
SECTION_HEADER_R = re.compile(r'^(\s*\[)([^\]]*)(\].*)$')
DIR_PIN_R = re.compile(r'^(\s*dir_pin\s*:\s*)([^#;]*)(.*)$', re.IGNORECASE)
AUTOSAVE_HEADER_R = re.compile(r'^#\*# \[([^\]]*)\]\s*$')
AUTOSAVE_OPTION_R = re.compile(r'^#\*# (\w+)\s*=\s*(.*)$')
INCLUDE_PREFIX = "include "


class AxisDiscovery:
    def __init__(self, config):
        self.printer = printer = config.get_printer()
        self.gcode = printer.lookup_object('gcode')
        self.wizard = None
        # Register g-code commands
        handlers = ['AXIS_DISCOVERY', 'AXIS_DISCOVERY_NEXT',
                    'AXIS_DISCOVERY_REPEAT', 'AXIS_DISCOVERY_ANSWER',
                    'AXIS_DISCOVERY_APPLY', 'AXIS_DISCOVERY_CANCEL']
        for cmd in handlers:
            func = getattr(self, 'cmd_' + cmd)
            desc = getattr(self, 'cmd_' + cmd + '_help', None)
            self.gcode.register_command(cmd, func, desc=desc)
        printer.register_event_handler("klippy:disconnect",
                                       self._handle_disconnect)

    def _handle_disconnect(self):
        self.wizard = None

    # Prompt emission (compatible with the web ui "action:prompt_*" protocol)
    def _prompt_action(self, action, msg=None):
        self.gcode.respond_raw("// action:prompt_%s%s"
                               % (action, "" if msg is None else " " + msg))

    def _prompt_begin(self, headline):
        self._prompt_action("begin", headline)

    def _prompt_text(self, msg):
        self._prompt_action("text", msg)

    def _prompt_button(self, msg):
        self._prompt_action("button", msg)

    def _prompt_footer_button(self, msg):
        self._prompt_action("footer_button", msg)

    def _prompt_show(self):
        self._prompt_action("show")

    def _prompt_end(self):
        self._prompt_action("end")

    # Wizard helpers
    def _check_not_printing(self):
        vsd = self.printer.lookup_object('virtual_sdcard', None)
        if vsd is not None and vsd.is_active():
            raise self.printer.command_error(
                "Command not available during a print")

    def _get_wizard(self, gcmd):
        if self.wizard is None:
            raise gcmd.error(
                "No axis discovery in progress (run AXIS_DISCOVERY first)")
        return self.wizard

    def _lookup_motors(self):
        toolhead = self.printer.lookup_object('toolhead')
        kin = toolhead.get_kinematics()
        if kin.__class__.__name__ != 'CartKinematics':
            raise self.printer.command_error(
                "AXIS_DISCOVERY supports [cartesian] kinematics only")
        rails = getattr(kin, 'rails', None)
        if rails is None or len(rails) != 3:
            raise self.printer.command_error(
                "AXIS_DISCOVERY is not supported with a dual carriage")
        raw_config = self.printer.lookup_object('configfile').get_status(
            0.)['config']
        motors = []
        for rail, axis in zip(rails, AXES):
            for mcu_stepper in rail.get_steppers():
                name = mcu_stepper.get_name()
                m = STEPPER_NAME_R.match(name)
                if m is None or m.group(1).lower() != axis:
                    raise self.printer.command_error(
                        "Unexpected stepper section '%s'" % (name,))
                dir_pin = raw_config.get(name, {}).get('dir_pin', '')
                motors.append({
                    'name': name, 'axis': axis, 'suffix': m.group(2),
                    'inverted': str(dir_pin).strip().startswith('!'),
                    'answer': None,
                })
        return motors

    def _move_current(self, wizard):
        self._check_not_printing()
        motor = wizard['motors'][wizard['index']]
        force_move = self.printer.lookup_object('force_move')
        stepper_enable = self.printer.lookup_object('stepper_enable')
        mcu_stepper = force_move.lookup_stepper(motor['name'])
        did_enable = stepper_enable.set_motors_enable([motor['name']], True)
        try:
            force_move.manual_move(mcu_stepper, wizard['distance'],
                                   wizard['speed'], wizard['accel'])
        finally:
            if did_enable:
                stepper_enable.set_motors_enable([motor['name']], False)

    def _show_question(self, wizard):
        motor = wizard['motors'][wizard['index']]
        self._prompt_begin("Axis Discovery (%d/%d)"
                           % (wizard['index'] + 1, len(wizard['motors'])))
        self._prompt_text("Moving %s in the positive direction..."
                          % (motor['name'],))
        self._prompt_text("Which carriage moved and in which direction?")
        for axis in AXES:
            self._prompt_action("button_group_start")
            self._prompt_button("%s+|AXIS_DISCOVERY_ANSWER I=%d AXIS=%s DIR=+"
                                % (axis.upper(), wizard['index'], axis))
            self._prompt_button("%s-|AXIS_DISCOVERY_ANSWER I=%d AXIS=%s DIR=-"
                                % (axis.upper(), wizard['index'], axis))
            self._prompt_action("button_group_end")
        self._prompt_footer_button("Move again|AXIS_DISCOVERY_REPEAT")
        self._prompt_footer_button("Cancel|AXIS_DISCOVERY_CANCEL")
        self._prompt_show()

    def _calc_changes(self, wizard):
        rename_map = {}
        invert_new = set()
        new_names = set()
        info = []
        for motor in wizard['motors']:
            axis, sign = motor['answer']
            new_name = "stepper_%s%s" % (axis, motor['suffix'])
            if new_name.lower() in new_names:
                raise self.printer.command_error(
                    "Inconsistent axis discovery results (duplicate motor"
                    " name '%s')" % (new_name,))
            new_names.add(new_name.lower())
            label = motor['name']
            if new_name != motor['name']:
                rename_map[motor['name'].lower()] = new_name.lower()
                label += " -> " + new_name
            if sign < 0:
                invert_new.add(new_name.lower())
                label += " (dir inverted)"
            info.append(label)
        return rename_map, invert_new, info

    def _show_summary(self, wizard):
        rename_map, invert_new, info = self._calc_changes(wizard)
        self._prompt_begin("Axis Discovery - confirm changes")
        if not rename_map and not invert_new:
            self._prompt_text("No changes required.")
        else:
            for line in info:
                self._prompt_text(line)
            self._prompt_text("The printer config is updated and the printer"
                              " restarts afterwards.")
            self._prompt_footer_button("Apply & Restart|AXIS_DISCOVERY_APPLY")
        self._prompt_footer_button("Cancel|AXIS_DISCOVERY_CANCEL")
        self._prompt_show()

    # G-code commands
    cmd_AXIS_DISCOVERY_help = "Interactively identify axis motors"
    def cmd_AXIS_DISCOVERY(self, gcmd):
        self._check_not_printing()
        distance = gcmd.get_float('DISTANCE', 10., above=0., maxval=100.)
        speed = gcmd.get_float('SPEED', 25., above=0., maxval=200.)
        accel = gcmd.get_float('ACCEL', 0., minval=0.)
        motors = self._lookup_motors()
        # Manual stepper moves invalidate the toolhead position
        toolhead = self.printer.lookup_object('toolhead')
        toolhead.get_kinematics().clear_homing_state(AXES)
        self.wizard = {'motors': motors, 'index': -1,
                       'distance': distance, 'speed': speed, 'accel': accel}
        self._prompt_begin("Axis Discovery")
        self._prompt_text("Detected motors:")
        for motor in motors:
            self._prompt_text("- %s (axis %s, dir_pin %s)"
                              % (motor['name'], motor['axis'].upper(),
                                 "inverted" if motor['inverted']
                                 else "normal"))
        self._prompt_text("Each motor is moved %.1fmm in the positive"
                          " direction. Watch which carriage moves and where"
                          " it goes." % (distance,))
        self._prompt_text("Homing is invalidated, the printer must be"
                          " re-homed after discovery.")
        self._prompt_footer_button("Start|AXIS_DISCOVERY_NEXT I=0")
        self._prompt_footer_button("Cancel|AXIS_DISCOVERY_CANCEL")
        self._prompt_show()

    cmd_AXIS_DISCOVERY_NEXT_help = "Move the next motor and request input"
    def cmd_AXIS_DISCOVERY_NEXT(self, gcmd):
        wizard = self._get_wizard(gcmd)
        expected = gcmd.get_int('I', None)
        if expected is not None and expected != wizard['index'] + 1:
            # Stale prompt button (eg, clicked twice); ignore it
            return
        if wizard['index'] + 1 >= len(wizard['motors']):
            self._show_summary(wizard)
            return
        wizard['index'] += 1
        self._move_current(wizard)
        self._show_question(wizard)

    cmd_AXIS_DISCOVERY_REPEAT_help = "Repeat the move of the current motor"
    def cmd_AXIS_DISCOVERY_REPEAT(self, gcmd):
        wizard = self._get_wizard(gcmd)
        if wizard['index'] < 0 or wizard['index'] >= len(wizard['motors']):
            raise gcmd.error("No motor is currently being identified")
        self._move_current(wizard)
        self._show_question(wizard)

    cmd_AXIS_DISCOVERY_ANSWER_help = "Report the observed axis motion"
    def cmd_AXIS_DISCOVERY_ANSWER(self, gcmd):
        wizard = self._get_wizard(gcmd)
        expected = gcmd.get_int('I', None)
        if expected is not None and expected != wizard['index']:
            raise gcmd.error(
                "Stale axis discovery answer (motor index %d expected)"
                % (wizard['index'],))
        if wizard['index'] < 0 or wizard['index'] >= len(wizard['motors']):
            raise gcmd.error("No motor is awaiting an answer")
        axis = gcmd.get('AXIS').lower()
        if axis not in AXES:
            raise gcmd.error("Invalid AXIS '%s'" % (axis,))
        direction = gcmd.get('DIR')
        if direction == '+':
            sign = 1
        elif direction == '-':
            sign = -1
        else:
            raise gcmd.error("Invalid DIR '%s' (must be + or -)"
                             % (direction,))
        motor = wizard['motors'][wizard['index']]
        for other in wizard['motors']:
            if other is motor or other['answer'] is None:
                continue
            if other['answer'][0] == axis:
                if other['axis'] != motor['axis']:
                    raise gcmd.error(
                        "Axis %s was already identified from %s; the answers"
                        " are inconsistent" % (axis.upper(), other['name']))
            elif other['axis'] == motor['axis']:
                raise gcmd.error(
                    "Motors of axis %s were observed on different axes"
                    " (%s and %s); the answers are inconsistent"
                    % (motor['axis'].upper(), other['name'], motor['name']))
        motor['answer'] = (axis, sign)
        wizard['index'] += 1
        if wizard['index'] < len(wizard['motors']):
            self._move_current(wizard)
            self._show_question(wizard)
        else:
            self._show_summary(wizard)

    cmd_AXIS_DISCOVERY_APPLY_help = "Apply axis discovery changes and restart"
    def cmd_AXIS_DISCOVERY_APPLY(self, gcmd):
        wizard = self._get_wizard(gcmd)
        for motor in wizard['motors']:
            if motor['answer'] is None:
                raise gcmd.error("Motor %s has not been identified yet"
                                 % (motor['name'],))
        self._check_not_printing()
        rename_map, invert_new, info = self._calc_changes(wizard)
        if not rename_map and not invert_new:
            self.wizard = None
            self._prompt_end()
            gcmd.respond_info("Axis discovery: no changes required")
            return
        self._apply_config(gcmd, rename_map, invert_new)
        self.wizard = None
        self._prompt_end()
        gcmd.respond_info("Axis discovery: config updated, restarting...")
        self.gcode.request_restart('restart')

    cmd_AXIS_DISCOVERY_CANCEL_help = "Abort the axis discovery wizard"
    def cmd_AXIS_DISCOVERY_CANCEL(self, gcmd):
        self.wizard = None
        self._prompt_end()
        gcmd.respond_info("Axis discovery cancelled (re-home the printer)")

    # Config file update support
    def _map_section_name(self, name, rename_map):
        stripped = name.strip()
        low = stripped.lower()
        if low in rename_map:
            new_name = rename_map[low]
        else:
            m = DEP_SECTION_R.match(low)
            if m is None or m.group(2) not in rename_map:
                return name
            new_name = "%s %s" % (m.group(1), rename_map[m.group(2)])
        return name.replace(stripped, new_name, 1)

    def _invert_pin(self, pin, section):
        pin = pin.strip()
        if not pin:
            raise self.printer.command_error(
                "Empty dir_pin in section '%s'" % (section,))
        if pin.startswith('!'):
            return pin[1:]
        return '!' + pin

    def _toggle_dir_pin(self, line_match, section):
        new_pin = self._invert_pin(line_match.group(2), section)
        return (line_match.group(1)
                + line_match.group(2).replace(
                    line_match.group(2).strip(), new_pin, 1)
                + line_match.group(3))

    def _toggle_autosave_dir_pin(self, option_match, section):
        new_pin = self._invert_pin(option_match.group(2), section)
        return "#*# %s = %s" % (option_match.group(1), new_pin)

    def _read_config_file(self, fname, files, visited):
        fname = os.path.abspath(fname)
        if fname in visited:
            raise self.printer.command_error(
                "Recursive include of config file '%s'" % (fname,))
        visited.add(fname)
        try:
            f = open(fname, 'r')
            data = f.read()
            f.close()
        except OSError:
            raise self.printer.command_error(
                "Unable to read config file '%s'" % (fname,))
        data = data.replace('\r\n', '\n')
        lines = data.split('\n')
        files[fname] = lines
        dirname = os.path.dirname(fname)
        for line in lines:
            hm = SECTION_HEADER_R.match(line)
            if hm is None:
                continue
            name = hm.group(2).strip()
            if not name.lower().startswith(INCLUDE_PREFIX):
                continue
            include_glob = os.path.join(dirname,
                                        name[len(INCLUDE_PREFIX):].strip())
            filenames = glob.glob(include_glob)
            if not filenames and not glob.has_magic(include_glob):
                raise self.printer.command_error(
                    "Include file '%s' does not exist" % (include_glob,))
            for include_fname in sorted(filenames):
                self._read_config_file(include_fname, files, visited)
        visited.remove(fname)

    def _apply_config(self, gcmd, rename_map, invert_new):
        cfgname = self.printer.get_start_args()['config_file']
        files = {}
        self._read_config_file(cfgname, files, set())
        # Apply renames and dir_pin inversions in memory (in the regular
        # config and in the "#*#" SAVE_CONFIG autosave block, which would
        # otherwise override the edits on the next start)
        toggle_counts = {}
        new_files = {}
        for fname, lines in files.items():
            cur_section = None
            in_autosave = False
            edits = 0
            out = []
            for line in lines:
                hm = SECTION_HEADER_R.match(line)
                am = None if hm is not None \
                    else AUTOSAVE_HEADER_R.match(line)
                if hm is not None:
                    name = hm.group(2).strip()
                    if name.lower().startswith(INCLUDE_PREFIX):
                        cur_section = None
                        out.append(line)
                        continue
                    new_name = self._map_section_name(name, rename_map)
                    if new_name != name:
                        line = (hm.group(1)
                                + hm.group(2).replace(name, new_name, 1)
                                + hm.group(3))
                        edits += 1
                    cur_section = new_name.strip().lower()
                    in_autosave = False
                    out.append(line)
                    continue
                if am is not None:
                    new_name = self._map_section_name(am.group(1), rename_map)
                    if new_name != am.group(1):
                        line = "#*# [%s]" % (new_name,)
                        edits += 1
                    cur_section = new_name.strip().lower()
                    in_autosave = True
                    out.append(line)
                    continue
                if cur_section in invert_new:
                    if in_autosave:
                        om = AUTOSAVE_OPTION_R.match(line)
                        if om is not None and om.group(1).lower() == 'dir_pin':
                            line = self._toggle_autosave_dir_pin(om,
                                                                 cur_section)
                            toggle_counts[cur_section] = \
                                toggle_counts.get(cur_section, 0) + 1
                            edits += 1
                    else:
                        dm = DIR_PIN_R.match(line)
                        if dm is not None:
                            line = self._toggle_dir_pin(dm, cur_section)
                            toggle_counts[cur_section] = \
                                toggle_counts.get(cur_section, 0) + 1
                            edits += 1
                out.append(line)
            if edits:
                new_files[fname] = out
        for name in invert_new:
            if name not in toggle_counts:
                raise gcmd.error("Unable to find dir_pin in section '%s'"
                                 % (name,))
        # Write all modified files to temporary names first so that a write
        # failure leaves the current config untouched
        stamp = time.strftime("-%Y%m%d_%H%M%S")
        pending = []
        try:
            for fname, lines in new_files.items():
                temp_name = fname + "_axisdiscovery.tmp"
                f = open(temp_name, 'w')
                f.write('\n'.join(lines))
                f.close()
                pending.append((fname, temp_name, fname + stamp))
        except os.error as e:
            for fname, temp_name, backup_name in pending:
                try:
                    os.remove(temp_name)
                except os.error:
                    pass
            raise gcmd.error("Unable to write config file: %s" % (e,))
        # Swap the new files into place (keeping a backup of each), rolling
        # back already-swapped files if a rename fails.  Note that the
        # backups intentionally do not keep a ".cfg" suffix (unlike
        # cmd_SAVE_CONFIG) so that common "[include *.cfg]" globs do not
        # load the pre-edit sections on the next start.
        committed = []
        try:
            for fname, temp_name, backup_name in pending:
                os.rename(fname, backup_name)
                os.rename(temp_name, fname)
                committed.append((fname, backup_name))
        except os.error as e:
            for fname, backup_name in reversed(committed):
                try:
                    os.rename(backup_name, fname)
                except os.error:
                    logging.exception(
                        "axis_discovery: Unable to restore '%s'" % (fname,))
            for fname, temp_name, backup_name in pending:
                try:
                    os.remove(temp_name)
                except os.error:
                    pass
            raise gcmd.error("Unable to write config file '%s': %s"
                             % (fname, e))
        for fname, backup_name in committed:
            logging.info("axis_discovery: wrote '%s' (backup in '%s')",
                         fname, backup_name)


def load_config(config):
    return AxisDiscovery(config)

Read 4 times, last 4 days ago