codekasten/gcode-tools/gcode-shifter/gcode_shift.py

#!/usr/bin/env python

import argparse
import os
import re
import sys
from decimal import Decimal as d

AXIS = "XYZ"
AXIS_REGEX = r"\b([%s][\+\-0-9\.]+)\b" % AXIS


class GCodeFilter(object):

    def __init__(self, lower, upper, func=None):
        self.pos = tuple([d(0)] * len(AXIS))
        self.lower = lower
        self.upper = upper
        self.recalc_func = func
        self.source_pos = PositionHandler()
        self.target_pos = PositionHandler()

    def is_inside(self, position):
        for i in range(len(position)):
            if not ((self.lower[i] is None or (self.lower[i] <= position[i])) and \
                    (self.upper[i] is None or (self.upper[i] >= position[i]))):
                return False
        return True

    def update_position(self, line):
        for comment_sep in ";(":
            if comment_sep in line:
                line = line.split(comment_sep)[0]
        new_pos = list(self.pos)
        for token in re.findall(AXIS_REGEX, line, flags=re.I):
            for i, axis in enumerate(AXIS):
                if token.upper().startswith(axis):
                    new_pos[i] = d(token[1:])
        new_pos = tuple(new_pos)
        if self.pos != new_pos:
            self.source_pos.update(new_pos)
            self.pos = new_pos
            return True
        else:
            return False

    def get_processed_lines(self, line):
        """ parse the line, update current position and replace line string
        """
        if self.update_position(line):
            result = []
            handler = LineHandler(line)
            for pos in self.transform_position():
                changed_axes = self.target_pos.get_changed_axes(pos)
                yield handler.get_line(pos, changed_axes)
                self.target_pos.update(pos)
        else:
            # no coordinate given
            yield line


class PositionHandler(object):

    def __init__(self):
        self.pos = tuple([0] * len(AXIS))
        self.pos_stack = [tuple(self.pos)]

    def push(self, position):
        self.pos_stack.insert(0, tuple(position))
        if len(self.pos_stack) > 10:
            self.pos_stack.pop(-1)

    def update(self, position):
        new_pos = tuple(position)
        if self.pos != new_pos:
            self.push(new_pos)
            self.pos = new_pos
            return True
        else:
            return False

    def get_changed_axes(self, position):
        for i, axis in enumerate(AXIS):
            if position[i] != self.pos[i]:
                yield (i, axis)


class LineHandler(object):

    def __init__(self, line):
        self._processed = []
        # remove trailing whitespace (and linebreak)
        self.line = line.rstrip()
        # all missing whitespace
        self._suffix = line[len(self.line):]
        self._axes = []

    def replace_value(self, match):
        token = match.group(0).upper()
        for i, axis in enumerate(AXIS):
            if token.startswith(axis):
                self._processed.append(i)
                if d(token[1:]) != self._position[i]:
                    return token[0] + str(self._position[i])
        else:
            return token

    def get_line(self, position, axes):
        result = []
        self._processed = []
        self._position = position
        self._axes = axes
        line = re.sub(AXIS_REGEX, self.replace_value, self.line, re.I)
        if not self._processed:
            # no axis definition
            line = self.line
        else:
            # at least one axis was defined
            for i, axis in axes:
                if not i in self._processed:
                    line += " %s%s" % (axis, str(position[i]))
        return line + self._suffix


class ShiftFilter(GCodeFilter):

    def __init__(self, *args, **kwargs):
        self.shift = kwargs.pop("shift", tuple([0] * len(AXIS)))
        super(ShiftFilter, self).__init__(*args, **kwargs)

    def transform_position(self):
        if self.is_inside(self.pos):
            yield [(axis + shift_axis)
                    for axis, shift_axis in zip(self.pos, self.shift)]
        else:
            yield self.pos


class DensifyFilter(GCodeFilter):

    def __init__(self, *args, **kwargs):
        self.loops = kwargs.pop("densify_loops")
        self.minimum_step = kwargs.pop("densify_minimum_step")
        self.direction = kwargs.pop("densify_dir").upper()
        self.direction_index = AXIS.index(self.direction)
        super(DensifyFilter, self).__init__(*args, **kwargs)

    def _densify_is_valid(self, p1, p2):
        for index, (v1, v2) in enumerate(zip(p1, p2)):
            if index == self.direction_index:
                if v1 == v2:
                    return False
                elif abs(v1 - v2) < self.minimum_step:
                    return False
                else:
                    pass
            elif v1 != v2:
                # all other direction vectors must be zero
                return False
            else:
                pass
        else:
            return True
    
    def _shift_point(self, pos, diff, reverse=False):
        if reverse:
            target = [p - d for p, d in zip(pos, diff)]
        else:
            target = [p + d for p, d in zip(pos, diff)]
        return tuple(target)

    def transform_position(self):
        stack = self.source_pos.pos_stack
        if (len(stack) > 2) and self.is_inside(stack[0]) and \
                self.is_inside(stack[1]) and \
                self._densify_is_valid(stack[0], stack[1]):
            # move sideways
            divider = d(float(self.loops * 2 + 1))
            dir_step = []
            for now, prev in zip(stack[0], stack[1]):
                value = (now - prev)/divider
                # adjust accuracy
                value = value.quantize(now)
                dir_step.append(value)
            # move forward and backward
            line_diff = [prev - prevprev
                    for prev, prevprev in zip(stack[1], stack[2])]
            current = stack[1]
            # run some loops
            for step in range(self.loops):
                # move to the side
                current = self._shift_point(current, dir_step)
                yield current
                # move back against line_diff
                current = self._shift_point(current, line_diff, reverse=True)
                yield current
                # move to the side
                current = self._shift_point(current, dir_step)
                yield current
                # move up against line_diff
                current = self._shift_point(current, line_diff, reverse=False)
                yield current
        # always move to the final point
        yield stack[0]


class CropFilter(GCodeFilter):

    def transform_position(self):
        is_inside = self.is_inside(self.pos)
        was_inside = self.pos_stack[1]
        result = None
        if is_inside and was_inside:
            # inside -> inside
            yield self.pos
        elif not is_inside and not was_inside:
            # outside -> outside
            pass
        else:
            # outside -> inside OR inside -> outside
            border_pos = self._recurse_border_position(self._previous_position,
                    pos, self.is_inside)
            # the border position is always the first step
            yield border_pos
            # omit the current position if we are outside
            if is_inside:
                yield self.pos

    def _recurse_border_position(self, p1, p2, depth_limit=20):
        """ simple and stupid: bisections between p1 and p2
        """
        p_middle = tuple([0.5 * (axis1 + axis2)
                for axis1, axis2 in zip(p1, p2)])
        if depth_limit < 0:
            return p_middle
        else:
            if self.is_inside(p_middle) == self.is_inside(p1):
                p1_new, p2_new = p_middle, p2
            else:
                p1_new, p2_new = p1, p_middle
            return self._recurse_border_position(p1_new, p2_new,
                    depth_limit=depth_limit-1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Shift parts of gcode")
    parser.add_argument('--minx', dest="minx", type=d)
    parser.add_argument('--maxx', dest="maxx", type=d)
    parser.add_argument('--miny', dest="miny", type=d)
    parser.add_argument('--maxy', dest="maxy", type=d)
    parser.add_argument('--minz', dest="minz", type=d)
    parser.add_argument('--maxz', dest="maxz", type=d)
    parser.add_argument('--shiftx', dest="shiftx", type=d, default=d(0))
    parser.add_argument('--shifty', dest="shifty", type=d, default=d(0))
    parser.add_argument('--shiftz', dest="shiftz", type=d, default=d(0))
    parser.add_argument('--rotatex', dest="rotatex", type=d, default=d(0))
    parser.add_argument('--rotatey', dest="rotatey", type=d, default=d(0))
    parser.add_argument('--rotatez', dest="rotatez", type=d, default=d(0))
    parser.add_argument('--densify-dir', dest="densify_dir",
            choices=("x", "y", "z"), default="x")
    parser.add_argument('--densify-loops', dest="densify_loops", type=int,
            default=1)
    parser.add_argument('--densify-minimum-step', dest="densify_minimum_step",
            type=d, default=d(0))
    parser.add_argument('action', \
            choices=("shift", "crop", "rotate", "densify"))
    infile = sys.stdin
    outfile = sys.stdout
    options = parser.parse_args()
    low = (options.minx, options.miny, options.minz)
    high = (options.maxx, options.maxy, options.maxz)
    if options.action == "shift":
        shift = (options.shiftx, options.shifty, options.shiftz)
        gcode_filter = ShiftFilter(low, high, shift=shift)
    elif options.action == "crop":
        gcode_filter = CropFilter(low, high)
    elif options.action == "rotate":
        rotate = (options.rotatex, options.rotatey, options.rotatez)
        gcode_filter = RotateFilter(low, high, rotate=rotate)
    elif options.action == "densify":
        gcode_filter = DensifyFilter(low, high, densify_dir=options.densify_dir,
                densify_loops=options.densify_loops,
                densify_minimum_step=options.densify_minimum_step)
    else:
        print >>sys.stderr, "No valid action choosen"
        sys.exit(1)
    for line in infile.readlines():
        for out_line in gcode_filter.get_processed_lines(line):
            outfile.write(out_line)