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

#!/usr/bin/env python

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

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


class GCodeFilter(object):

    def __init__(self, destination, 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()
        self.destination = destination
        self.destination_buffer = collections.deque()
        self._status = {}

    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 parse(self, line):
        """ parse the line, update current position and replace line string
        """
        if self.update_position(line):
            result = []
            handler = LineHandler(line)
            positions = list(self.transform_position())
            if not positions:
                self.destination_buffer.append(handler.get_line(None, None))
            else:
                for pos in positions:
                    if callable(pos):
                        pos, line = pos()
                        self.destination_buffer.append(line)
                    else:
                        changed_axes = self.target_pos.get_changed_axes(pos)
                        self.destination_buffer.append(handler.get_line(pos,
                                changed_axes))
                    self.target_pos.update(pos)
        else:
            # no coordinate given
            self.destination_buffer.append(line)
        self.write()

    def _parse_rapid(self, line):
        # return True / False / None for G0/G1/other commands
        for item in line.upper().split():
            if item == "G0":
                return True
            if item == "G1":
                return False
        return None

    def _parse_rapid_noop(self, line):
        rapid = self._parse_rapid(line)
        if rapid is None:
            return False
        else:
            # only one item (G0 or G1)
            return len(line.strip().split()) == 1

    def write(self):
        """ add all no-op filters here (e.g. "GO / G1" -> "G1")
        """
        if not self.destination:
            raise IOError("GCodeFilter was closed before")
        pending_rapid = None
        while self.destination_buffer:
            item = self.destination_buffer.popleft()
            if item is None:
                # skip invalid items
                continue
            rapid = self._parse_rapid(item)
            rapid_noop = self._parse_rapid_noop(item)
            if rapid_noop:
                if rapid == self._status.get("rapid"):
                    # skip this no-op (and any pending ones)
                    pending_rapid = None
                    continue
                else:
                    pending_rapid = item
                    continue
            else:
                if (not pending_rapid is None) and (rapid is None) and \
                        (self._parse_rapid(pending_rapid) != rapid):
                    self.destination.write(pending_rapid)
                pending_rapid = None
            # no matches -> write it
            self.destination.write(item)
            if not rapid is None:
                self._status["rapid"] = rapid
        if not pending_rapid is None:
            # put it back to the list
            self.destination_buffer.appendleft(pending_rapid)

    def close(self):
        if self.destination is None:
            return
        self.write()
        # write all pending items
        while self.destination_buffer:
            self.destination.write(self.destination_buffer.popleft())
        self.destination.close()
        self.destination = None

    def __del__(self):
        self.close()


class PositionHandler(object):

    def __init__(self):
        self.pos = tuple([None] * len(AXIS))
        self.stack = [tuple(self.pos)]

    def push(self, position):
        self.stack.insert(0, tuple(position))
        if len(self.stack) > 10:
            self.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 self.pos[i] is None:
                continue
            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):
        if position is None:
            # remove all coordinates from the line
            # (e.g. for a modal switch ("G0 X..") outside of a crop space)
            line = re.sub(AXIS_REGEX, "", self.line, re.I)
            for i in range(10):
                line = line.replace("  ", " ")
            if line.strip():
                return line.rstrip() + self._suffix
            else:
                # invalid line - skip it
                return None
        else:
            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 MatrixFilter(GCodeFilter):

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

    def _multiply_with_matrix(self, pos):
        rows = []
        for row in self.matrix:
            rows.append(sum([(p * r).quantize(p) for p, r in zip(pos, row)]))
        return tuple(rows)

    def transform_position(self):
        if self.is_inside(self.pos):
            yield self._multiply_with_matrix(self.pos)
        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.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):
        stack = self.source_pos.stack
        target_stack = self.target_pos.stack
        is_inside = self.is_inside(stack[0])
        was_inside = (len(stack) < 2) or self.is_inside(stack[1])
        result = None
        if is_inside and was_inside:
            # inside -> inside
            yield stack[0]
        elif not is_inside and not was_inside:
            # outside -> outside
            # only apply relative upward moves (safety height)
            z_index = AXIS.lower().index("z")
            diff = [(now - prev) for now, prev in zip(stack[0], stack[1])]
            for index, value in enumerate(diff):
                if index == z_index:
                    if value <= 0:
                        break
                elif value != 0:
                    break
            else:
                # this looks like an upward move
                if target_stack[0][z_index] <= stack[0][z_index]:
                    new_target = list(target_stack[0])
                    new_target[z_index] = stack[0][z_index]
                    yield tuple(new_target)
        else:
            # outside -> inside OR inside -> outside
            if len(stack) >= 2:
                border_pos = self._recurse_border_position(stack[1], stack[0])
                # adapt the precision to the precision of the input
                border_pos = tuple([value.quantize(template)
                        for value, template in zip(border_pos, stack[0])])
                # the border position is always the first step
                if is_inside:
                    # outside -> inside
                    line = "G0 %s" % " ".join([axis.upper() + str(value)
                                for axis, value in zip(AXIS, border_pos)])
                    yield lambda: (border_pos, line + os.linesep)
                else:
                    # inside -> outside
                    yield border_pos
            # omit the current position if we are outside
            if is_inside:
                yield stack[0]

    def _recurse_border_position(self, p1, p2, depth_limit=20):
        """ simple and stupid: bisections between p1 and p2
        """
        p_middle = tuple([d(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)


def get_rotate_matrix(axis, angle):
    rot = [d(0)] * 3
    rot[AXIS.lower().index(axis.lower())] = d(1)
    # clockwise angle
    angle_pi = -angle * d(math.pi / 180.0)
    sin = d(math.sin(angle_pi))
    cos = d(math.cos(angle_pi))
    return ((cos + rot[0] * rot[0] * (1 - cos),
                rot[0] * rot[1] * (1-cos) - rot[2] * sin,
                rot[0] * rot[2] * (1-cos) + rot[1] * sin),
            (rot[1] * rot[0] *(1 - cos) + rot[2] * sin,
                cos + rot[1] * rot[1] * (1 - cos),
                rot[1] * rot[2] * (1 - cos) - rot[0] * sin),
            (rot[2] * rot[0] * (1 - cos) - rot[1] * sin,
                rot[2] * rot[1] * (1 - cos) + rot[0] * sin,
                cos + rot[2] * rot[2] * (1 - cos)))


if __name__ == "__main__":
    axes_choices = tuple(AXIS.lower() + AXIS.upper())
    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('--rotate-axis', dest="rotate_axis",
            choices=axes_choices, default=AXIS[0])
    parser.add_argument('--rotate-angle', dest="rotate_angle", type=d,
            default=d(90))
    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)
    common_args = (outfile, low, high)
    if options.action == "shift":
        shift = (options.shiftx, options.shifty, options.shiftz)
        gcode_filter = ShiftFilter(*common_args, shift=shift)
    elif options.action == "crop":
        gcode_filter = CropFilter(*common_args)
    elif options.action == "rotate":
        matrix = get_rotate_matrix(options.rotate_axis, options.rotate_angle)
        gcode_filter = MatrixFilter(*common_args, matrix=matrix)
    elif options.action == "densify":
        gcode_filter = DensifyFilter(*common_args,
                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():
        gcode_filter.parse(line)