<commit_before>import importlib
import re
import token
import tokenize
from io import StringIO
from pathlib import Path
from collections import defaultdict, deque, namedtuple, ChainMap

TOK_COMMENT = 'comment'
TOK_TEXT = 'text'
TOK_VAR = 'var'
TOK_BLOCK = 'block'

tag_re = re.compile(r'{%\s*(?P<block>.+?)\s*%}|{{\s*(?P<var>.+?)\s*}}|{#\s*(?P<comment>.+?)\s*#}', re.DOTALL)
Token = namedtuple('Token', 'type content')

def tokenise(template):
    upto = 0
    for m in tag_re.finditer(template):
        start, end = m.span()
        if upto < start:
            yield Token(TOK_TEXT, template[upto:start])
        upto = end
        mode = m.lastgroup
        yield Token(mode, m[mode].strip())
    if upto < len(template):
        yield Token(TOK_TEXT, template[upto:])

class TemplateLoader(dict):
    def __init__(self, paths):
        self.paths = [Path(path).resolve() for path in paths]

    def load(self, name, encoding='utf8'):
        for path in self.paths:
            full_path = path / name
            if full_path.is_file():
                return Template(full_path.read_text(encoding), loader=self, name=name)
        raise LookupError(name)

    def __missing__(self, key):
        self[key] = tmpl = self.load(key)
        return tmpl

class Context(ChainMap):
    def __init__(self, *args):
        super().__init__(*args)
        self.maps.append({'True': True, 'False': False, 'None': None})

    def push(self, data=None):
        self.maps.insert(0, data or {})
        return self

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        self.maps.pop(0)

class Nodelist(list):
    def render(self, context, output):
        for node in self:
            node.render(context, output)

    def nodes_by_type(self, node_type):
        for node in self:
            if isinstance(node, node_type):
                yield node
            if isinstance(node, BlockNode):
                yield from node.nodes_by_type(node_type)

class Template:
    def __init__(self, src, loader=None, name=None):
        self.tokens, self.loader = tokenise(src), loader
        self.nodelist = self.parse_nodelist([])
        self.name = name  # So we can report where the fault was

    def parse(self):
        for tok in self.tokens:
            if tok.type == TOK_TEXT:
                yield TextTag(tok.content)
            elif tok.type == TOK_VAR:
                yield VarTag(tok.content)
            elif tok.type == TOK_BLOCK:
                m = re.match(r'\w+', tok.content)
                if not m:
                    raise SyntaxError(tok)
                yield BlockNode.__tags__[m.group(0)].parse(tok.content[m.end(0):].strip(), self)

    def parse_nodelist(self, ends):
        nodelist = Nodelist()
        try:
            node = next(self.parse())
            while node.name not in ends:
                nodelist.append(node)
                node = next(self.parse())
        except StopIteration:
            node = None
        nodelist.endnode = node
        return nodelist

    def render(self, context, output=None):
        if not isinstance(context, Context):
            context = Context(context)
        if output is None:
            dest = StringIO()
        else:
            dest = output
        self.nodelist.render(context, dest)
        if output is None:
            return dest.getvalue()

class AstLiteral:
    def __init__(self, arg):
        self.arg = arg

    def resolve(self, context):
        return self.arg

class AstContext:
    def __init__(self, arg):
        self.arg = arg

    def resolve(self, context):
        return context.get(self.arg, '')

class AstLookup:
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def resolve(self, context):
        left = self.left.resolve(context)
        right = self.right.resolve(context)

        return left[right]

class AstAttr:
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def resolve(self, context):
        left = self.left.resolve(context)

        return getattr(left, self.right)

class AstCall:
    def __init__(self, func):
        self.func = func
        self.args = []

    def add_arg(self, arg):
        self.args.append(arg)

    def resolve(self, context):
        func = self.func.resolve(context)
        args = [arg.resolve(context) for arg in self.args]

        return func(*args)

class Expression:
    def __init__(self, source):
        self.tokens = tokenize.generate_tokens(StringIO(source).readline)
        self.next()  # prime the first token

    def next(self):
        self.current = next(self.tokens)
        return self.current

    @staticmethod
    def parse(s):
        p = Expression(s)
        result = p._parse()

        if p.current.exact_type not in (token.NEWLINE, token.ENDMARKER):
            raise SyntaxError(f'Parse ended unexpectedly: {p.current}')

        return result

    def parse_kwargs(self):
        kwargs = {}

        tok = self.current

        while tok.exact_type != token.ENDMARKER:
            if tok == token.NEWLINE:
                tok = self.next()
                continue

            if tok.exact_type != token.NAME:
                raise SyntaxError(f'Expected name, found {tok}')
            name = tok.string
            tok = self.next()

            if tok.exact_type != token.EQUAL:
                raise SyntaxError(f'Expected =, found {tok}')
            tok = self.next()

            kwargs[name] = self._parse()

            tok = self.next()

        return kwargs

    def _parse(self):
        tok = self.current

        if tok.exact_type in (token.ENDMARKER, token.COMMA):
            return  # TODO

        if tok.exact_type == token.STRING:
            self.next()
            return AstLiteral(tok.string[1:-1])

        if tok.exact_type == token.NUMBER:
            self.next()
            try:
                value = int(tok.string)
            except ValueError:
                value = float(tok.string)
            return AstLiteral(value)

        if tok.exact_type == token.NAME:
            state = AstContext(tok.string)

            while True:
                tok = self.next()

                if tok.exact_type == token.DOT:
                    tok = self.next()
                    if tok.exact_type != token.NAME:
                        raise SyntaxError(f"Invalid attr lookup: {tok}")
                    state = AstAttr(state, tok.string)

                elif tok.exact_type == token.LSQB:
                    self.next()
                    right = self._parse()
                    state = AstLookup(state, right)
                    if self.current.exact_type != token.RSQB:
                        raise SyntaxError(f"Expected ] but found {self.current}")

                elif tok.exact_type == token.LPAR:
                    state = AstCall(state)
                    self.next()
                    while self.current.exact_type != token.RPAR:
                        arg = self._parse()
                        state.add_arg(arg)
                        if self.current.exact_type != token.COMMA:
                            break
                        self.next()

                    if self.current.exact_type != token.RPAR:
                        raise SyntaxError(f"Expected ( but found {self.current}")

                    self.next()

                else:
                    break

            return state

        raise SyntaxError(f'Unexpected token: {tok}')

class Node:
    name = None

    def __init__(self, content):
        self.content = content

    def render(self, context, output):
        pass  # pragma: no cover

class TextTag(Node):
    def render(self, context, output):
        output.write(self.content)

class VarTag(Node):
    def __init__(self, content):
        self.expr = Expression.parse(content)

    def render(self, context, output):
        output.write(str(self.expr.resolve(context)))

class BlockNode(Node):
    __tags__ = {}
    child_nodelists = ('nodelist',)

    def __init_subclass__(cls, *, name):
        super().__init_subclass__()
        cls.name = name
        BlockNode.__tags__[name] = cls
        return cls

    @classmethod
    def parse(cls, content, parser):
        return cls(content)

    def nodes_by_type(self, node_type):
        for attr in self.child_nodelists:
            nodelist = getattr(self, attr, None)
            if nodelist:
                yield from nodelist.nodes_by_type(node_type)

class ForTag(BlockNode, name='for'):
    child_nodelists = ('nodelist', 'elselist')

    def __init__(self, argname, iterable, nodelist, elselist):
        self.argname, self.iterable, self.nodelist, self.elselist = argname, iterable, nodelist, elselist

    @classmethod
    def parse(cls, content, parser):
        argname, iterable = content.split(' in ', 1)
        nodelist = parser.parse_nodelist({'endfor', 'else'})
        elselist = parser.parse_nodelist({'endfor'}) if nodelist.endnode.name == 'else' else None
        return cls(argname.strip(), Expression.parse(iterable.strip()), nodelist, elselist)

    def render(self, context, output):
        iterable = self.iterable.resolve(context)
        if self.elselist and not iterable:
            self.elselist.render(context, output)
        else:
            with context.push():
                for idx, item in enumerate(iterable):
                    context.update({'loopcounter': idx, self.argname: item})
                    self.nodelist.render(context, output)

class ElseTag(BlockNode, name='else'):
    pass

class EndforTag(BlockNode, name='endfor'):
    pass

class IfTag(BlockNode, name='if'):
    child_nodelists = ('nodelist', 'elselist')

    def __init__(self, condition, nodelist, elselist):
        condition, inv = re.subn(r'^not\s+', '', condition, count=1)
        self.inv, self.condition = bool(inv), Expression.parse(condition)
        self.nodelist, self.elselist = nodelist, elselist

    @classmethod
    def parse(cls, content, parser):
        nodelist = parser.parse_nodelist({'endif', 'else'})
        elselist = parser.parse_nodelist({'endif'}) if nodelist.endnode.name == 'else' else None
        return cls(content, nodelist, elselist)

    def render(self, context, output):
        if self.test_condition(context):
            self.nodelist.render(context, output)
        elif self.elselist:
            self.elselist.render(context, output)

    def test_condition(self, context):
        return self.inv ^ bool(self.condition.resolve(context))

class EndifTag(BlockNode, name='endif'):
    pass

class IncludeTag(BlockNode, name='include'):

    def __init__(self, template_name, kwargs, loader):
        self.template_name, self.kwargs, self.loader = template_name, kwargs, loader

    @classmethod
    def parse(cls, content, parser):
        if parser.loader is None:
            raise RuntimeError("Can't use {% include %} without a bound Loader")
        tokens = Expression(content)
        template_name = tokens._parse()
        kwargs = tokens.parse_kwargs()
        return cls(template_name, kwargs, parser.loader)

    def render(self, context, output):
        name = self.template_name.resolve(context)
        tmpl = self.loader[name]
        kwargs = {key: expr.resolve(context) for key, expr in self.kwargs.items()}
        with context.push(kwargs):
            tmpl.render(context, output)

class LoadTag(BlockNode, name='load'):

    @classmethod
    def parse(cls, content, parser):
        importlib.import_module(content)
        return cls(None)

class ExtendsTag(BlockNode, name='extends'):

    def __init__(self, parent, loader, nodelist):
        self.parent, self.loader, self.nodelist = parent, loader, nodelist

    @classmethod
    def parse(cls, content, parser):
        parent = Expression.parse(content)
        nodelist = parser.parse_nodelist([])
        return cls(parent, parser.loader, nodelist)

    def render(self, context, output):
        parent = self.loader[self.parent.resolve(context)]
        block_context = getattr(context, 'block_context', None)
        new = False
        if block_context is None:
            block_context = context.block_context = defaultdict(deque)
            new = True
        for block in self.nodelist.nodes_by_type(BlockTag):
            block_context[block.block_name].append(block)
        if parent.nodelist[0].name != 'extends':
            for block in parent.nodelist.nodes_by_type(BlockTag):
                block_context[block.block_name].append(block)
        parent.render(context, output)
        if new:
            del context.block_context

class BlockTag(BlockNode, name='block'):

    def __init__(self, name, nodelist):
        self.block_name, self.nodelist = name, nodelist

    @classmethod
    def parse(cls, content, parser):
        name = re.match(r'\w+', content).group(0)
        nodelist = parser.parse_nodelist({'endblock'})
        return cls(name, nodelist)

    def render(self, context, output):
        block_context = getattr(context, 'block_context', None)
        if not block_context:
            block = self
        else:
            block = block_context[self.block_name].popleft()
        with context.push():
            block.nodelist.render(context, output)

class EndBlockTag(BlockNode, name='endblock'):
    pass

class WithTag(BlockNode, name='with'):

    def __init__(self, kwargs, nodelist):
        self.kwargs, self.nodelist = kwargs, nodelist

    @classmethod
    def parse(cls, content, parser):
        kwargs = Expression(content).parse_kwargs()
        nodelist = parser.parse_nodelist({'endwith'})
        return cls(kwargs, nodelist)

    def render(self, context, output):
        kwargs = {key: value.resolve(context) for key, value in self.kwargs.items()}
        with context.push(kwargs):
            self.nodelist.render(context, output)

class EndWithTag(BlockNode, name='endwith'):
    pass

class CaseTag(BlockNode, name='case'):

    def __init__(self, term, nodelist):
        self.term, self.nodelist = term, nodelist

    @classmethod
    def parse(cls, content, parser):
        term = Expression.parse(content)
        nodelist = parser.parse_nodelist(['endcase'])
        else_found = False
        for node in nodelist:
            if node.name not in {'when', 'else'}:
                raise SyntaxError(f"Only 'when' and 'else' allowed as children of case. Found: {node}")
            if node.name == 'else':
                if else_found:
                    raise SyntaxError('Case tag can only have one else child')
                else_found = True
        nodelist.sort(key=lambda x: x.name, reverse=True)
        return cls(term, nodelist)

    def render(self, context, output):
        value = self.term.resolve(context)
        for node in self.nodelist:
            if node.name == 'when':
                other = node.term.resolve(context)
            else:
                other = value
            if value == other:
                node.render(context, output)
                return

class WhenTag(BlockNode, name='when'):

    def __init__(self, term, nodelist):
        self.term, self.nodelist = term, nodelist

    @classmethod
    def parse(cls, content, parser):
        term = Expression.parse(content)
        nodelist = parser.parse_nodelist()
        return cls(term, nodelist)

    def render(self, context, output):
        self.nodelist.render(context, output)

class EndCaseTag(BlockNode, name='endcase'):
    pass
<commit_msg>Report parsing errors more meaningfully<commit_after>import importlib
import re
import token
import tokenize
from io import StringIO
from pathlib import Path
from collections import defaultdict, deque, namedtuple, ChainMap

TOK_COMMENT = 'comment'
TOK_TEXT = 'text'
TOK_VAR = 'var'
TOK_BLOCK = 'block'

tag_re = re.compile(r'{%\s*(?P<block>.+?)\s*%}|{{\s*(?P<var>.+?)\s*}}|{#\s*(?P<comment>.+?)\s*#}', re.DOTALL)
Token = namedtuple('Token', 'type content')

def tokenise(template):
    upto = 0
    for m in tag_re.finditer(template):
        start, end = m.span()
        if upto < start:
            yield Token(TOK_TEXT, template[upto:start])
        upto = end
        mode = m.lastgroup
        yield Token(mode, m[mode].strip())
    if upto < len(template):
        yield Token(TOK_TEXT, template[upto:])

class TemplateLoader(dict):
    def __init__(self, paths):
        self.paths = [Path(path).resolve() for path in paths]

    def load(self, name, encoding='utf8'):
        for path in self.paths:
            full_path = path / name
            if full_path.is_file():
                return Template(full_path.read_text(encoding), loader=self, name=name)
        raise LookupError(name)

    def __missing__(self, key):
        self[key] = tmpl = self.load(key)
        return tmpl

class Context(ChainMap):
    def __init__(self, *args):
        super().__init__(*args)
        self.maps.append({'True': True, 'False': False, 'None': None})

    def push(self, data=None):
        self.maps.insert(0, data or {})
        return self

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        self.maps.pop(0)

class Nodelist(list):
    def render(self, context, output):
        for node in self:
            node.render(context, output)

    def nodes_by_type(self, node_type):
        for node in self:
            if isinstance(node, node_type):
                yield node
            if isinstance(node, BlockNode):
                yield from node.nodes_by_type(node_type)

class Template:
    def __init__(self, src, loader=None, name=None):
        self.tokens, self.loader = tokenise(src), loader
        self.name = name  # So we can report where the fault was
        self.nodelist = self.parse_nodelist([])

    def parse(self):
      try:
        for tok in self.tokens:
            if tok.type == TOK_TEXT:
                yield TextTag(tok.content)
            elif tok.type == TOK_VAR:
                yield VarTag(tok.content)
            elif tok.type == TOK_BLOCK:
                m = re.match(r'\w+', tok.content)
                if not m:
                    raise SyntaxError(tok)
                yield BlockNode.__tags__[m.group(0)].parse(tok.content[m.end(0):].strip(), self)
      except Exception as ex:
        raise RuntimeError(f"Error parsing template {self.name}") from ex

    def parse_nodelist(self, ends):
        nodelist = Nodelist()
        try:
            node = next(self.parse())
            while node.name not in ends:
                nodelist.append(node)
                node = next(self.parse())
        except StopIteration:
            node = None
        nodelist.endnode = node
        return nodelist

    def render(self, context, output=None):
        if not isinstance(context, Context):
            context = Context(context)
        if output is None:
            dest = StringIO()
        else:
            dest = output
        self.nodelist.render(context, dest)
        if output is None:
            return dest.getvalue()

class AstLiteral:
    def __init__(self, arg):
        self.arg = arg

    def resolve(self, context):
        return self.arg

class AstContext:
    def __init__(self, arg):
        self.arg = arg

    def resolve(self, context):
        return context.get(self.arg, '')

class AstLookup:
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def resolve(self, context):
        left = self.left.resolve(context)
        right = self.right.resolve(context)

        return left[right]

class AstAttr:
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def resolve(self, context):
        left = self.left.resolve(context)

        return getattr(left, self.right)

class AstCall:
    def __init__(self, func):
        self.func = func
        self.args = []

    def add_arg(self, arg):
        self.args.append(arg)

    def resolve(self, context):
        func = self.func.resolve(context)
        args = [arg.resolve(context) for arg in self.args]

        return func(*args)

class Expression:
    def __init__(self, source):
        self.tokens = tokenize.generate_tokens(StringIO(source).readline)
        self.next()  # prime the first token

    def next(self):
        self.current = next(self.tokens)
        return self.current

    @staticmethod
    def parse(s):
        p = Expression(s)
        result = p._parse()

        if p.current.exact_type not in (token.NEWLINE, token.ENDMARKER):
            raise SyntaxError(f'Parse ended unexpectedly: {p.current}')

        return result

    def parse_kwargs(self):
        kwargs = {}

        tok = self.current

        while tok.exact_type != token.ENDMARKER:
            if tok == token.NEWLINE:
                tok = self.next()
                continue

            if tok.exact_type != token.NAME:
                raise SyntaxError(f'Expected name, found {tok}')
            name = tok.string
            tok = self.next()

            if tok.exact_type != token.EQUAL:
                raise SyntaxError(f'Expected =, found {tok}')
            tok = self.next()

            kwargs[name] = self._parse()

            tok = self.next()

        return kwargs

    def _parse(self):
        tok = self.current

        if tok.exact_type in (token.ENDMARKER, token.COMMA):
            return  # TODO

        if tok.exact_type == token.STRING:
            self.next()
            return AstLiteral(tok.string[1:-1])

        if tok.exact_type == token.NUMBER:
            self.next()
            try:
                value = int(tok.string)
            except ValueError:
                value = float(tok.string)
            return AstLiteral(value)

        if tok.exact_type == token.NAME:
            state = AstContext(tok.string)

            while True:
                tok = self.next()

                if tok.exact_type == token.DOT:
                    tok = self.next()
                    if tok.exact_type != token.NAME:
                        raise SyntaxError(f"Invalid attr lookup: {tok}")
                    state = AstAttr(state, tok.string)

                elif tok.exact_type == token.LSQB:
                    self.next()
                    right = self._parse()
                    state = AstLookup(state, right)
                    if self.current.exact_type != token.RSQB:
                        raise SyntaxError(f"Expected ] but found {self.current}")

                elif tok.exact_type == token.LPAR:
                    state = AstCall(state)
                    self.next()
                    while self.current.exact_type != token.RPAR:
                        arg = self._parse()
                        state.add_arg(arg)
                        if self.current.exact_type != token.COMMA:
                            break
                        self.next()

                    if self.current.exact_type != token.RPAR:
                        raise SyntaxError(f"Expected ( but found {self.current}")

                    self.next()

                else:
                    break

            return state

        raise SyntaxError(f'Unexpected token: {tok}')

class Node:
    name = None

    def __init__(self, content):
        self.content = content

    def render(self, context, output):
        pass  # pragma: no cover

class TextTag(Node):
    def render(self, context, output):
        output.write(self.content)

class VarTag(Node):
    def __init__(self, content):
        self.expr = Expression.parse(content)

    def render(self, context, output):
        output.write(str(self.expr.resolve(context)))

class BlockNode(Node):
    __tags__ = {}
    child_nodelists = ('nodelist',)

    def __init_subclass__(cls, *, name):
        super().__init_subclass__()
        cls.name = name
        BlockNode.__tags__[name] = cls
        return cls

    @classmethod
    def parse(cls, content, parser):
        return cls(content)

    def nodes_by_type(self, node_type):
        for attr in self.child_nodelists:
            nodelist = getattr(self, attr, None)
            if nodelist:
                yield from nodelist.nodes_by_type(node_type)

class ForTag(BlockNode, name='for'):
    child_nodelists = ('nodelist', 'elselist')

    def __init__(self, argname, iterable, nodelist, elselist):
        self.argname, self.iterable, self.nodelist, self.elselist = argname, iterable, nodelist, elselist

    @classmethod
    def parse(cls, content, parser):
        argname, iterable = content.split(' in ', 1)
        nodelist = parser.parse_nodelist({'endfor', 'else'})
        elselist = parser.parse_nodelist({'endfor'}) if nodelist.endnode.name == 'else' else None
        return cls(argname.strip(), Expression.parse(iterable.strip()), nodelist, elselist)

    def render(self, context, output):
        iterable = self.iterable.resolve(context)
        if self.elselist and not iterable:
            self.elselist.render(context, output)
        else:
            with context.push():
                for idx, item in enumerate(iterable):
                    context.update({'loopcounter': idx, self.argname: item})
                    self.nodelist.render(context, output)

class ElseTag(BlockNode, name='else'):
    pass

class EndforTag(BlockNode, name='endfor'):
    pass

class IfTag(BlockNode, name='if'):
    child_nodelists = ('nodelist', 'elselist')

    def __init__(self, condition, nodelist, elselist):
        condition, inv = re.subn(r'^not\s+', '', condition, count=1)
        self.inv, self.condition = bool(inv), Expression.parse(condition)
        self.nodelist, self.elselist = nodelist, elselist

    @classmethod
    def parse(cls, content, parser):
        nodelist = parser.parse_nodelist({'endif', 'else'})
        elselist = parser.parse_nodelist({'endif'}) if nodelist.endnode.name == 'else' else None
        return cls(content, nodelist, elselist)

    def render(self, context, output):
        if self.test_condition(context):
            self.nodelist.render(context, output)
        elif self.elselist:
            self.elselist.render(context, output)

    def test_condition(self, context):
        return self.inv ^ bool(self.condition.resolve(context))

class EndifTag(BlockNode, name='endif'):
    pass

class IncludeTag(BlockNode, name='include'):

    def __init__(self, template_name, kwargs, loader):
        self.template_name, self.kwargs, self.loader = template_name, kwargs, loader

    @classmethod
    def parse(cls, content, parser):
        if parser.loader is None:
            raise RuntimeError("Can't use {% include %} without a bound Loader")
        tokens = Expression(content)
        template_name = tokens._parse()
        kwargs = tokens.parse_kwargs()
        return cls(template_name, kwargs, parser.loader)

    def render(self, context, output):
        name = self.template_name.resolve(context)
        tmpl = self.loader[name]
        kwargs = {key: expr.resolve(context) for key, expr in self.kwargs.items()}
        with context.push(kwargs):
            tmpl.render(context, output)

class LoadTag(BlockNode, name='load'):

    @classmethod
    def parse(cls, content, parser):
        importlib.import_module(content)
        return cls(None)

class ExtendsTag(BlockNode, name='extends'):

    def __init__(self, parent, loader, nodelist):
        self.parent, self.loader, self.nodelist = parent, loader, nodelist

    @classmethod
    def parse(cls, content, parser):
        parent = Expression.parse(content)
        nodelist = parser.parse_nodelist([])
        return cls(parent, parser.loader, nodelist)

    def render(self, context, output):
        parent = self.loader[self.parent.resolve(context)]
        block_context = getattr(context, 'block_context', None)
        new = False
        if block_context is None:
            block_context = context.block_context = defaultdict(deque)
            new = True
        for block in self.nodelist.nodes_by_type(BlockTag):
            block_context[block.block_name].append(block)
        if parent.nodelist[0].name != 'extends':
            for block in parent.nodelist.nodes_by_type(BlockTag):
                block_context[block.block_name].append(block)
        parent.render(context, output)
        if new:
            del context.block_context

class BlockTag(BlockNode, name='block'):

    def __init__(self, name, nodelist):
        self.block_name, self.nodelist = name, nodelist

    @classmethod
    def parse(cls, content, parser):
        name = re.match(r'\w+', content).group(0)
        nodelist = parser.parse_nodelist({'endblock'})
        return cls(name, nodelist)

    def render(self, context, output):
        block_context = getattr(context, 'block_context', None)
        if not block_context:
            block = self
        else:
            block = block_context[self.block_name].popleft()
        with context.push():
            block.nodelist.render(context, output)

class EndBlockTag(BlockNode, name='endblock'):
    pass

class WithTag(BlockNode, name='with'):

    def __init__(self, kwargs, nodelist):
        self.kwargs, self.nodelist = kwargs, nodelist

    @classmethod
    def parse(cls, content, parser):
        kwargs = Expression(content).parse_kwargs()
        nodelist = parser.parse_nodelist({'endwith'})
        return cls(kwargs, nodelist)

    def render(self, context, output):
        kwargs = {key: value.resolve(context) for key, value in self.kwargs.items()}
        with context.push(kwargs):
            self.nodelist.render(context, output)

class EndWithTag(BlockNode, name='endwith'):
    pass

class CaseTag(BlockNode, name='case'):

    def __init__(self, term, nodelist):
        self.term, self.nodelist = term, nodelist

    @classmethod
    def parse(cls, content, parser):
        term = Expression.parse(content)
        nodelist = parser.parse_nodelist(['endcase'])
        else_found = False
        for node in nodelist:
            if node.name not in {'when', 'else'}:
                raise SyntaxError(f"Only 'when' and 'else' allowed as children of case. Found: {node}")
            if node.name == 'else':
                if else_found:
                    raise SyntaxError('Case tag can only have one else child')
                else_found = True
        nodelist.sort(key=lambda x: x.name, reverse=True)
        return cls(term, nodelist)

    def render(self, context, output):
        value = self.term.resolve(context)
        for node in self.nodelist:
            if node.name == 'when':
                other = node.term.resolve(context)
            else:
                other = value
            if value == other:
                node.render(context, output)
                return

class WhenTag(BlockNode, name='when'):

    def __init__(self, term, nodelist):
        self.term, self.nodelist = term, nodelist

    @classmethod
    def parse(cls, content, parser):
        term = Expression.parse(content)
        nodelist = parser.parse_nodelist()
        return cls(term, nodelist)

    def render(self, context, output):
        self.nodelist.render(context, output)

class EndCaseTag(BlockNode, name='endcase'):
    pass
<|endoftext|>


<commit_before>#!/usr/bin/python
#pin 11,10,6,5 motore 1
#pin 14,13,12,3 motore 2

import wiringpi2 as wiringpi
import time
import threading
from subprocess import call

class movement(threading.Thread):
	accel = 10
	realspeed = 0.1
	speedwheel = 0.1
	decremento = 0.1
	go = False
	started = False
	avanti = False
	def run(self):
		while self.started:
			if self.go:
				self.decremento = (self.realspeed-self.speedwheel)/self.accel
				wiringpi.digitalWrite(self.motor[self.ix],1)
				self.px = self.ix
				if self.avanti:
					self.ix +=1
				else:
					self.ix -=1
				if (self.ix == 4):
					self.ix = 0
				if (self.ix == -1):
					self.ix = 3
				time.sleep(self.realspeed)
				wiringpi.digitalWrite(self.motor[self.ix],1)
				time.sleep(self.realspeed)
				wiringpi.digitalWrite(self.motor[self.px],0)
				if self.realspeed > self.speedwheel:
					self.realspeed -= self.decremento
				else:
					self.realspeed = self.speedwheel
			else:
				for x in self.motor:
					wiringpi.digitalWrite(x,0)
				self.realspeed = 0.1

wiringpi.wiringPiSetup()
wiringpi.pinMode(11,1)
wiringpi.pinMode(10,1)
wiringpi.pinMode(6,1)
wiringpi.pinMode(5,1)
wiringpi.pinMode(14,1)
wiringpi.pinMode(13,1)
wiringpi.pinMode(12,1)
wiringpi.pinMode(3,1)

stepper1 = movement()
stepper2 = movement()

stepper1.motor = [10,11,6,5]
stepper2.motor = [14,13,12,3]
stepper1.go = False
stepper2.go = False
stepper1.started = True
stepper2.started = True
stepper1.start()
stepper2.start()
rw = 1.0
lw = 1.0
direction =""
denominator = 1.0
#call(["sudo killall servod"],shell=True)
#call(["sudo ./servod"],shell=True)

while (direction !="quit"):
	for x in stepper1.motor:
		wiringpi.digitalWrite(x,0)
	for x in stepper2.motor:
		wiringpi.digitalWrite(x,0)
	direction = raw_input("action?(fd/bk/lt/rt/ad/quit) ")
	if direction != "ad" and direction != "fd" and direction != "bk" and direction !="lt" and direction !="rt" and direction !="quit":
		print "only fd,bk,lt,rt,quit"
		continue
	if direction != "quit":
		if direction !="ad":
			try:
				denominator = input("speed?(number, 0 exit) ")
			except NameError:
				print "only numbers"
				continue
			try:
				duration = input("duration?(seconds) ")
			except NameError:
				print "only numbers"
				continue
		if direction == "bk":
			stepper1.avanti = True
			stepper2.avanti = True
			rw = lw = 1.0/denominator
		elif direction == "fd":
			stepper1.avanti = False
			stepper2.avanti = False
			rw = lw = 1.0/denominator
		elif direction == "rt":
			stepper1.avanti = True
			stepper2.avanti = False
			rw = lw = 1.0/denominator
		elif direction == "lt":
			stepper1.avanti = False
			stepper2.avanti = True
			rw = lw = 1.0/denominator
#		elif direction == "pu":
#			call(["echo 0=100 > /dev/servoblaster"],shell=True)
#			duration = 0
#		elif direction == "pd":
#			call(["echo 0=80 > /dev/servoblaster"],shell=True)
#			duration = 0
		elif direction == "ad":
			try:
				rw = input("right wheel interval? ")
				dir = input("right wheel direction? (0/1)")
				if dir == 1:
					stepper1.avanti = True
				elif dir == 0:
					stepper1.avanti = False
				else:
					print "0/1 only"
					continue
				lw = input("left wheel interval?  ")
				dir = input("left wheel direction?  (0/1)")
				if dir == 1:
					stepper2.avanti = True
				elif dir == 0:
					stepper2.avanti = False
				else:
					print "0/1 only"
					continue
				duration = input("duration? ")
			except NameError:
				print "numbers only"
				continue
		time.sleep (0.5)
		stepper1.ix = 0
		stepper2.ix = 0
		stepper1.px = 0
		stepper2.px = 0
		stepper1.speedwheel = rw
		stepper2.speedwheel = lw
		stepper1.go = True
		stepper2.go = True
		time.sleep(duration)
		stepper1.go = False
		stepper2.go = False
	else:
		print "exit"
		call(["sudo killall servod"],shell=True)
		stepper1.started = False
		stepper2.started = False
<commit_msg>aggiunto controllo penna<commit_after>#!/usr/bin/python
#pin 11,10,6,5 motore 1
#pin 14,13,12,3 motore 2

import wiringpi2 as wiringpi
import time
import threading
from subprocess import call

class movement(threading.Thread):
#	accel = 10
	realspeed = 0.1
	speedwheel = 0.1
#	decremento = 0.1
	go = False
	started = False
	avanti = False
	def run(self):
		while self.started:
			if self.go:
#				self.decremento = (self.realspeed-self.speedwheel)/self.accel
				wiringpi.digitalWrite(self.motor[self.ix],1)
				self.px = self.ix
				if self.avanti:
					self.ix +=1
				else:
					self.ix -=1
				if (self.ix == 4):
					self.ix = 0
				if (self.ix == -1):
					self.ix = 3
				time.sleep(self.realspeed)
				wiringpi.digitalWrite(self.motor[self.ix],1)
				time.sleep(self.realspeed)
				wiringpi.digitalWrite(self.motor[self.px],0)
				if self.realspeed > self.speedwheel:
					#self.realspeed -= self.decremento
					self.realspeed -= 0.01
				else:
					self.realspeed = self.speedwheel
			else:
				for x in self.motor:
					wiringpi.digitalWrite(x,0)
				self.realspeed = 0.1

wiringpi.wiringPiSetup()
wiringpi.pinMode(11,1)
wiringpi.pinMode(10,1)
wiringpi.pinMode(6,1)
wiringpi.pinMode(5,1)
wiringpi.pinMode(14,1)
wiringpi.pinMode(13,1)
wiringpi.pinMode(12,1)
wiringpi.pinMode(3,1)

stepper1 = movement()
stepper2 = movement()

stepper1.motor = [10,11,6,5]
stepper2.motor = [14,13,12,3]
stepper1.go = False
stepper2.go = False
stepper1.started = True
stepper2.started = True
stepper1.start()
stepper2.start()
rw = 1.0
lw = 1.0
direction =""
denominator = 1.0
call(["sudo killall servod"],shell=True)
call(["sudo ./servod"],shell=True)

while (direction !="quit"):
	for x in stepper1.motor:
		wiringpi.digitalWrite(x,0)
	for x in stepper2.motor:
		wiringpi.digitalWrite(x,0)
	direction = raw_input("action?(fd/bk/lt/rt/pu/pd/ad/quit) ")
	if direction != "ad" and direction != "fd" and direction != "bk" and direction !="lt" and direction !="rt" and direction !="quit":
		print "only fd,bk,lt,rt,quit"
		continue
	if direction != "quit":
		if direction !="ad":
			try:
				denominator = input("speed?(number, 0 exit) ")
			except NameError:
				print "only numbers"
				continue
			try:
				duration = input("duration?(seconds) ")
			except NameError:
				print "only numbers"
				continue
		if direction == "bk":
			stepper1.avanti = True
			stepper2.avanti = True
			rw = lw = 1.0/denominator
		elif direction == "fd":
			stepper1.avanti = False
			stepper2.avanti = False
			rw = lw = 1.0/denominator
		elif direction == "rt":
			stepper1.avanti = True
			stepper2.avanti = False
			rw = lw = 1.0/denominator
		elif direction == "lt":
			stepper1.avanti = False
			stepper2.avanti = True
			rw = lw = 1.0/denominator
		elif direction == "pu":
			call(["echo 0=170 > /dev/servoblaster"],shell=True)
			duration = 0
		elif direction == "pd":
			call(["echo 0=185 > /dev/servoblaster"],shell=True)
			duration = 0
		elif direction == "ad":
			try:
				rw = input("right wheel interval? ")
				dir = input("right wheel direction? (0/1)")
				if dir == 1:
					stepper1.avanti = True
				elif dir == 0:
					stepper1.avanti = False
				else:
					print "0/1 only"
					continue
				lw = input("left wheel interval?  ")
				dir = input("left wheel direction?  (0/1)")
				if dir == 1:
					stepper2.avanti = True
				elif dir == 0:
					stepper2.avanti = False
				else:
					print "0/1 only"
					continue
				duration = input("duration? ")
			except NameError:
				print "numbers only"
				continue
		time.sleep (0.5)
		stepper1.ix = 0
		stepper2.ix = 0
		stepper1.px = 0
		stepper2.px = 0
		stepper1.speedwheel = rw
		stepper2.speedwheel = lw
		stepper1.go = True
		stepper2.go = True
		time.sleep(duration)
		stepper1.go = False
		stepper2.go = False
	else:
		print "exit"
		call(["sudo killall servod"],shell=True)
		stepper1.started = False
		stepper2.started = False
<|endoftext|>


<commit_before>import SimpleITK as sitk
import numpy, pywt, logging
from itertools import chain

logger = logging.getLogger(__name__)

def getHistogram(binwidth, parameterValues):
  # Start binning form the first value lesser than or equal to the minimum value and evenly dividable by binwidth
  lowBound = min(parameterValues) - (min(parameterValues) % binwidth)
  # Add + binwidth to ensure the maximum value is included in the range generated by numpu.arange
  highBound = max(parameterValues) + binwidth

  binedges = numpy.arange(lowBound, highBound, binwidth)

  binedges[-1] += 1 # ensures that max(self.targertVoxelArray) is binned to upper bin by numpy.digitize

  return numpy.histogram(parameterValues, bins=binedges)

def binImage(binwidth, parameterValues, parameterMatrix = None, parameterMatrixCoordinates = None):
  histogram = getHistogram(binwidth, parameterValues)
  parameterMatrix[parameterMatrixCoordinates] = numpy.digitize(parameterValues, histogram[1])

  return parameterMatrix, histogram

def generateAngles(size, maxDistance= 1):
  """
  Generate all possible angles from distance 1 until maxDistance in 3D.
  E.g. for d = 1, 13 angles are generated (representing the 26-connected region).
  For d = 2, 13 + 49 = 62 angles are generated (representing the 26 connected region for distance 1, and the 98
  connected region for distance 2)

  Impossible angles (where 'neighbouring' voxels will always be outside delineation) are deleted.

  :param size: dimensions (z, x, y) of the bounding box of the tumor mask.
  :param maxDistance: [1] Maximum distance between center voxel and neighbour
  :return: numpy array with shape (N, 3), where N is the number of unique angles
  """

  angles = []

  for z in xrange(1, maxDistance+1):
    angles.append((0, 0, z))
    for y in xrange(-maxDistance, maxDistance+1):
      angles.append((0, z, y))
      for x in xrange(-maxDistance, maxDistance+1):
        angles.append((z, y, x))

  angles = numpy.array(angles)

  angles = numpy.delete(angles, numpy.where(numpy.min(size - numpy.abs(angles), 1) <= 0), 0)

  return angles

def cropToTumorMask(imageNode, maskNode, label=1):
  """
  Create a sitkImage of the segmented region of the image based on the input label.

  Create a sitkImage of the labelled region of the image, cropped to have a
  cuboid shape equal to the ijk boundaries of the label.

  Returns both the cropped version of the image and the cropped version of the labelmap.
  """
  global logger

  oldMaskID = maskNode.GetPixelID()
  maskNode = sitk.Cast(maskNode, sitk.sitkInt32)
  size = numpy.array(maskNode.GetSize())

  #Determine bounds
  lsif = sitk.LabelStatisticsImageFilter()
  lsif.Execute(imageNode, maskNode)
  bb = numpy.array(lsif.GetBoundingBox(label))

  ijkMinBounds = bb[0::2]
  ijkMaxBounds = size - bb[1::2] - 1

  #Crop Image
  logger.debug('Cropping to size %s', (bb[1::2] - bb[0::2]) + 1)
  cif = sitk.CropImageFilter()
  cif.SetLowerBoundaryCropSize(ijkMinBounds)
  cif.SetUpperBoundaryCropSize(ijkMaxBounds)
  croppedImageNode = cif.Execute(imageNode)
  croppedMaskNode = cif.Execute(maskNode)

  croppedMaskNode = sitk.Cast(croppedMaskNode, oldMaskID)

  return croppedImageNode, croppedMaskNode

def resampleImage(imageNode, maskNode, resampledPixelSpacing, interpolator=sitk.sitkBSpline, label=1, padDistance=5):
  """Resamples image or label to the specified pixel spacing (The default interpolator is Bspline)

  'imageNode' is a SimpleITK Object, and 'resampledPixelSpacing' is the output pixel spacing.
  Enumerator references for interpolator:
  0 - sitkNearestNeighbor
  1 - sitkLinear
  2 - sitkBSpline
  3 - sitkGaussian
  """
  global logger

  if imageNode is None or maskNode is None:
    return None

  oldSpacing = numpy.array(imageNode.GetSpacing())

  # If current spacing is equal to resampledPixelSpacing, no interpolation is needed,
  # crop/pad image using cropTumorMaskToCube
  if numpy.array_equal(oldSpacing, resampledPixelSpacing):
    return cropToTumorMask(imageNode, maskNode)

  # Determine bounds of cropped volume in terms of original Index coordinate space
  lssif = sitk.LabelShapeStatisticsImageFilter()
  lssif.Execute(maskNode)
  bb = numpy.array(lssif.GetBoundingBox(label))  # LBound and size of the bounding box, as (L_X, L_Y, L_Z, S_X, S_Y, S_Z)

  # Do not resample in those directions where labelmap spans only one slice.
  oldSize = bb[3:]
  resampledPixelSpacing = numpy.where(oldSize != 1, resampledPixelSpacing, oldSpacing)

  spacingRatio = oldSpacing / resampledPixelSpacing

  # Determine bounds of cropped volume in terms of new Index coordinate space,
  # round down for lowerbound and up for upperbound to ensure entire segmentation is captured (prevent data loss)
  # Pad with an extra .5 to prevent data loss in case of upsampling. For Ubound this is (-1 + 0.5 = -0.5)
  bbNewLBound = numpy.floor( (bb[:3] - 0.5) * spacingRatio - padDistance)
  bbNewUBound = numpy.ceil( (bb[:3] + bb[3:] - 0.5) * spacingRatio + padDistance)

  # Ensure resampling is not performed outside bounds of original image
  maxUbound = numpy.ceil(numpy.array(imageNode.GetSize()) * spacingRatio) - 1
  bbNewLBound = numpy.where(bbNewLBound < 0, 0, bbNewLBound)
  bbNewUBound = numpy.where(bbNewUBound > maxUbound, maxUbound, bbNewUBound)

  # Calculate the new size. Cast to int to prevent error in sitk.
  newSize = numpy.array(bbNewUBound - bbNewLBound + 1, dtype='int')

  # Determine continuous index of bbNewLBound in terms of the original Index coordinate space
  bbOriginalLBound = bbNewLBound / spacingRatio

  # Origin is located in center of first voxel, e.g. 1/2 of the spacing
  # from Corner, which corresponds to 0 in the original Index coordinate space.
  # The new spacing will be in 0 the new Index coordinate space. Here we use continuous
  # index to calculate where the new 0 of the new Index coordinate space (of the original volume
  # in terms of the original spacing, and add the minimum bounds of the cropped area to
  # get the new Index coordinate space of the cropped volume in terms of the original Index coordinate space.
  # Then use the ITK functionality to bring the contiuous index into the physical space (mm)
  newOriginIndex = numpy.array(.5*(resampledPixelSpacing-oldSpacing)/oldSpacing)
  newCroppedOriginIndex = newOriginIndex + bbOriginalLBound
  newOrigin = imageNode.TransformContinuousIndexToPhysicalPoint(newCroppedOriginIndex)

  oldImagePixelType = imageNode.GetPixelID()
  oldMaskPixelType = maskNode.GetPixelID()

  imageDirection = numpy.array(imageNode.GetDirection())

  logger.debug('Applying resampling (spacing %s and size %s)', resampledPixelSpacing, newSize)

  rif = sitk.ResampleImageFilter()

  rif.SetOutputSpacing(resampledPixelSpacing)
  rif.SetOutputDirection(imageDirection)
  rif.SetSize(newSize)
  rif.SetOutputOrigin(newOrigin)

  rif.SetOutputPixelType(oldImagePixelType)
  rif.SetInterpolator(interpolator)
  resampledImageNode = rif.Execute(imageNode)

  rif.SetOutputPixelType(oldMaskPixelType)
  rif.SetInterpolator(sitk.sitkNearestNeighbor)
  resampledMaskNode = rif.Execute(maskNode)

  return resampledImageNode,resampledMaskNode

#
# Use the SimpleITK LaplacianRecursiveGaussianImageFilter
# on the input image with the given sigmaValue and return
# the filtered image.
# If sigmaValue is not greater than zero, return the input image.
#
def applyLoG(inputImage, sigmaValue=0.5):
  global logger
  if sigmaValue > 0.0:
    size = numpy.array(inputImage.GetSize())
    spacing = numpy.array(inputImage.GetSpacing())
    if numpy.all(size >= numpy.ceil(sigmaValue / spacing) + 1):
      lrgif = sitk.LaplacianRecursiveGaussianImageFilter()
      lrgif.SetNormalizeAcrossScale(True)
      lrgif.SetSigma(sigmaValue)
      return lrgif.Execute(inputImage)
    else:
      logger.warning('applyLoG: sigma/spacing + 1 must be greater than the size of the inputImage: %g', sigmaValue)
      return None
  else:
    logger.warning('applyLoG: sigma must be greater than 0.0: %g', sigmaValue)
    return None

def applyThreshold(inputImage, lowerThreshold, upperThreshold, insideValue=None, outsideValue=0):
  # this mode is useful to generate the mask of thresholded voxels
  if insideValue:
    tif = sitk.BinaryThresholdImageFilter()
    tif.SetInsideValue(insideValue)
    tif.SetLowerThreshold(lowerThreshold)
    tif.SetUpperThreshold(upperThreshold)
  else:
    tif = sitk.ThresholdImageFilter()
    tif.SetLower(lowerThreshold)
    tif.SetUpper(upperThreshold)
  tif.SetOutsideValue(outsideValue)
  return tif.Execute(inputImage)

def swt3(inputImage, wavelet="coif1", level=1, start_level=0):
  matrix = sitk.GetArrayFromImage(inputImage)
  matrix = numpy.asarray(matrix)
  data = matrix.copy()
  if data.ndim != 3:
    raise ValueError("Expected 3D data array")

  original_shape = matrix.shape
  adjusted_shape = tuple([dim+1 if dim % 2 != 0 else dim for dim in original_shape])
  data = numpy.resize(data, adjusted_shape)

  if not isinstance(wavelet, pywt.Wavelet):
    wavelet = pywt.Wavelet(wavelet)

  for i in range(0, start_level):
    H, L = decompose_i(data, wavelet)
    LH, LL = decompose_j(L, wavelet)
    LLH, LLL = decompose_k(LL, wavelet)

    data = LLL.copy()

  ret = []
  for i in range(start_level, start_level + level):
    H, L = decompose_i(data, wavelet)

    HH, HL = decompose_j(H, wavelet)
    LH, LL = decompose_j(L, wavelet)

    HHH, HHL = decompose_k(HH, wavelet)
    HLH, HLL = decompose_k(HL, wavelet)
    LHH, LHL = decompose_k(LH, wavelet)
    LLH, LLL = decompose_k(LL, wavelet)

    data = LLL.copy()

    dec = {'HHH': HHH,
           'HHL': HHL,
           'HLH': HLH,
           'HLL': HLL,
           'LHH': LHH,
           'LHL': LHL,
           'LLH': LLH}
    for decName, decImage in dec.iteritems():
      decTemp = decImage.copy()
      decTemp= numpy.resize(decTemp, original_shape)
      sitkImage = sitk.GetImageFromArray(decTemp)
      sitkImage.CopyInformation(inputImage)
      dec[decName] = sitkImage

    ret.append(dec)

  data= numpy.resize(data, original_shape)
  approximation = sitk.GetImageFromArray(data)
  approximation.CopyInformation(inputImage)

  return approximation, ret

def decompose_i(data, wavelet):
  #process in i:
  H, L = [], []
  i_arrays = chain.from_iterable(numpy.transpose(data,(0,1,2)))
  for i_array in i_arrays:
    cA, cD = pywt.swt(i_array, wavelet, level=1, start_level=0)[0]
    H.append(cD)
    L.append(cA)
  H = numpy.hstack(H).reshape(data.shape)
  L = numpy.hstack(L).reshape(data.shape)
  return H, L

def decompose_j(data, wavelet):
  #process in j:
  H, L = [], []
  j_arrays = chain.from_iterable(numpy.transpose(data,(0,1,2)))
  for j_array in j_arrays:
    cA, cD = pywt.swt(j_array, wavelet, level=1, start_level=0)[0]
    H.append(cD)
    L.append(cA)
  H = numpy.asarray( [slice.T for slice in numpy.split(numpy.vstack(H), data.shape[0])] )
  L = numpy.asarray( [slice.T for slice in numpy.split(numpy.vstack(L), data.shape[0])] )
  return H, L

def decompose_k(data, wavelet):
  #process in k:
  H, L = [], []
  k_arrays = chain.from_iterable(numpy.transpose(data,(1,2,0)))
  for k_array in k_arrays:
    cA, cD = pywt.swt(k_array, wavelet, level=1, start_level=0)[0]
    H.append(cD)
    L.append(cA)
  H = numpy.dstack(H).reshape(data.shape)
  L = numpy.dstack(L).reshape(data.shape)
  return H, L

def applySquare(inputImage):
  sif = sitk.SquareImageFilter()
  im = sif.Execute(inputImage)
  return _scaleToOriginalRange(inputImage, im)

def applySquareRoot(inputImage):
  srif = sitk.SqrtImageFilter()
  im = srif.Execute(inputImage)
  return _scaleToOriginalRange(inputImage, im)

def applyLogarithm(inputImage):
  lif = sitk.LogImageFilter()
  im = lif.Execute(inputImage)
  return _scaleToOriginalRange(inputImage, im)

def applyExponential(inputImage):
  eif = sitk.ExpImageFilter()
  im = eif.Execute(inputImage)
  return _scaleToOriginalRange(inputImage, im)

def _scaleToOriginalRange(originalImage, filteredImage):
  mmif = sitk.MinimumMaximumImageFilter()
  mmif.Execute(originalImage)
  im_max = mmif.GetMaximum()
  mmif.Execute(filteredImage)
  fl_max = mmif.GetMaximum()

  ssif = sitk.ShiftScaleImageFilter()
  ssif.SetScale(im_max / fl_max)

  im = ssif.Execute(filteredImage)
  im.CopyInformation(originalImage)
  return im
<commit_msg>BUG: Revert to numpy functionality for image filters<commit_after>import SimpleITK as sitk
import numpy, pywt, logging
from itertools import chain

logger = logging.getLogger(__name__)

def getHistogram(binwidth, parameterValues):
  # Start binning form the first value lesser than or equal to the minimum value and evenly dividable by binwidth
  lowBound = min(parameterValues) - (min(parameterValues) % binwidth)
  # Add + binwidth to ensure the maximum value is included in the range generated by numpu.arange
  highBound = max(parameterValues) + binwidth

  binedges = numpy.arange(lowBound, highBound, binwidth)

  binedges[-1] += 1 # ensures that max(self.targertVoxelArray) is binned to upper bin by numpy.digitize

  return numpy.histogram(parameterValues, bins=binedges)

def binImage(binwidth, parameterValues, parameterMatrix = None, parameterMatrixCoordinates = None):
  histogram = getHistogram(binwidth, parameterValues)
  parameterMatrix[parameterMatrixCoordinates] = numpy.digitize(parameterValues, histogram[1])

  return parameterMatrix, histogram

def generateAngles(size, maxDistance= 1):
  """
  Generate all possible angles from distance 1 until maxDistance in 3D.
  E.g. for d = 1, 13 angles are generated (representing the 26-connected region).
  For d = 2, 13 + 49 = 62 angles are generated (representing the 26 connected region for distance 1, and the 98
  connected region for distance 2)

  Impossible angles (where 'neighbouring' voxels will always be outside delineation) are deleted.

  :param size: dimensions (z, x, y) of the bounding box of the tumor mask.
  :param maxDistance: [1] Maximum distance between center voxel and neighbour
  :return: numpy array with shape (N, 3), where N is the number of unique angles
  """

  angles = []

  for z in xrange(1, maxDistance+1):
    angles.append((0, 0, z))
    for y in xrange(-maxDistance, maxDistance+1):
      angles.append((0, z, y))
      for x in xrange(-maxDistance, maxDistance+1):
        angles.append((z, y, x))

  angles = numpy.array(angles)

  angles = numpy.delete(angles, numpy.where(numpy.min(size - numpy.abs(angles), 1) <= 0), 0)

  return angles

def cropToTumorMask(imageNode, maskNode, label=1):
  """
  Create a sitkImage of the segmented region of the image based on the input label.

  Create a sitkImage of the labelled region of the image, cropped to have a
  cuboid shape equal to the ijk boundaries of the label.

  Returns both the cropped version of the image and the cropped version of the labelmap.
  """
  global logger

  oldMaskID = maskNode.GetPixelID()
  maskNode = sitk.Cast(maskNode, sitk.sitkInt32)
  size = numpy.array(maskNode.GetSize())

  #Determine bounds
  lsif = sitk.LabelStatisticsImageFilter()
  lsif.Execute(imageNode, maskNode)
  bb = numpy.array(lsif.GetBoundingBox(label))

  ijkMinBounds = bb[0::2]
  ijkMaxBounds = size - bb[1::2] - 1

  #Crop Image
  logger.debug('Cropping to size %s', (bb[1::2] - bb[0::2]) + 1)
  cif = sitk.CropImageFilter()
  cif.SetLowerBoundaryCropSize(ijkMinBounds)
  cif.SetUpperBoundaryCropSize(ijkMaxBounds)
  croppedImageNode = cif.Execute(imageNode)
  croppedMaskNode = cif.Execute(maskNode)

  croppedMaskNode = sitk.Cast(croppedMaskNode, oldMaskID)

  return croppedImageNode, croppedMaskNode

def resampleImage(imageNode, maskNode, resampledPixelSpacing, interpolator=sitk.sitkBSpline, label=1, padDistance=5):
  """Resamples image or label to the specified pixel spacing (The default interpolator is Bspline)

  'imageNode' is a SimpleITK Object, and 'resampledPixelSpacing' is the output pixel spacing.
  Enumerator references for interpolator:
  0 - sitkNearestNeighbor
  1 - sitkLinear
  2 - sitkBSpline
  3 - sitkGaussian
  """
  global logger

  if imageNode is None or maskNode is None:
    return None

  oldSpacing = numpy.array(imageNode.GetSpacing())

  # If current spacing is equal to resampledPixelSpacing, no interpolation is needed,
  # crop/pad image using cropTumorMaskToCube
  if numpy.array_equal(oldSpacing, resampledPixelSpacing):
    return cropToTumorMask(imageNode, maskNode)

  # Determine bounds of cropped volume in terms of original Index coordinate space
  lssif = sitk.LabelShapeStatisticsImageFilter()
  lssif.Execute(maskNode)
  bb = numpy.array(lssif.GetBoundingBox(label))  # LBound and size of the bounding box, as (L_X, L_Y, L_Z, S_X, S_Y, S_Z)

  # Do not resample in those directions where labelmap spans only one slice.
  oldSize = bb[3:]
  resampledPixelSpacing = numpy.where(oldSize != 1, resampledPixelSpacing, oldSpacing)

  spacingRatio = oldSpacing / resampledPixelSpacing

  # Determine bounds of cropped volume in terms of new Index coordinate space,
  # round down for lowerbound and up for upperbound to ensure entire segmentation is captured (prevent data loss)
  # Pad with an extra .5 to prevent data loss in case of upsampling. For Ubound this is (-1 + 0.5 = -0.5)
  bbNewLBound = numpy.floor( (bb[:3] - 0.5) * spacingRatio - padDistance)
  bbNewUBound = numpy.ceil( (bb[:3] + bb[3:] - 0.5) * spacingRatio + padDistance)

  # Ensure resampling is not performed outside bounds of original image
  maxUbound = numpy.ceil(numpy.array(imageNode.GetSize()) * spacingRatio) - 1
  bbNewLBound = numpy.where(bbNewLBound < 0, 0, bbNewLBound)
  bbNewUBound = numpy.where(bbNewUBound > maxUbound, maxUbound, bbNewUBound)

  # Calculate the new size. Cast to int to prevent error in sitk.
  newSize = numpy.array(bbNewUBound - bbNewLBound + 1, dtype='int')

  # Determine continuous index of bbNewLBound in terms of the original Index coordinate space
  bbOriginalLBound = bbNewLBound / spacingRatio

  # Origin is located in center of first voxel, e.g. 1/2 of the spacing
  # from Corner, which corresponds to 0 in the original Index coordinate space.
  # The new spacing will be in 0 the new Index coordinate space. Here we use continuous
  # index to calculate where the new 0 of the new Index coordinate space (of the original volume
  # in terms of the original spacing, and add the minimum bounds of the cropped area to
  # get the new Index coordinate space of the cropped volume in terms of the original Index coordinate space.
  # Then use the ITK functionality to bring the contiuous index into the physical space (mm)
  newOriginIndex = numpy.array(.5*(resampledPixelSpacing-oldSpacing)/oldSpacing)
  newCroppedOriginIndex = newOriginIndex + bbOriginalLBound
  newOrigin = imageNode.TransformContinuousIndexToPhysicalPoint(newCroppedOriginIndex)

  oldImagePixelType = imageNode.GetPixelID()
  oldMaskPixelType = maskNode.GetPixelID()

  imageDirection = numpy.array(imageNode.GetDirection())

  logger.debug('Applying resampling (spacing %s and size %s)', resampledPixelSpacing, newSize)

  rif = sitk.ResampleImageFilter()

  rif.SetOutputSpacing(resampledPixelSpacing)
  rif.SetOutputDirection(imageDirection)
  rif.SetSize(newSize)
  rif.SetOutputOrigin(newOrigin)

  rif.SetOutputPixelType(oldImagePixelType)
  rif.SetInterpolator(interpolator)
  resampledImageNode = rif.Execute(imageNode)

  rif.SetOutputPixelType(oldMaskPixelType)
  rif.SetInterpolator(sitk.sitkNearestNeighbor)
  resampledMaskNode = rif.Execute(maskNode)

  return resampledImageNode,resampledMaskNode

#
# Use the SimpleITK LaplacianRecursiveGaussianImageFilter
# on the input image with the given sigmaValue and return
# the filtered image.
# If sigmaValue is not greater than zero, return the input image.
#
def applyLoG(inputImage, sigmaValue=0.5):
  global logger
  if sigmaValue > 0.0:
    size = numpy.array(inputImage.GetSize())
    spacing = numpy.array(inputImage.GetSpacing())
    if numpy.all(size >= numpy.ceil(sigmaValue / spacing) + 1):
      lrgif = sitk.LaplacianRecursiveGaussianImageFilter()
      lrgif.SetNormalizeAcrossScale(True)
      lrgif.SetSigma(sigmaValue)
      return lrgif.Execute(inputImage)
    else:
      logger.warning('applyLoG: sigma/spacing + 1 must be greater than the size of the inputImage: %g', sigmaValue)
      return None
  else:
    logger.warning('applyLoG: sigma must be greater than 0.0: %g', sigmaValue)
    return None

def applyThreshold(inputImage, lowerThreshold, upperThreshold, insideValue=None, outsideValue=0):
  # this mode is useful to generate the mask of thresholded voxels
  if insideValue:
    tif = sitk.BinaryThresholdImageFilter()
    tif.SetInsideValue(insideValue)
    tif.SetLowerThreshold(lowerThreshold)
    tif.SetUpperThreshold(upperThreshold)
  else:
    tif = sitk.ThresholdImageFilter()
    tif.SetLower(lowerThreshold)
    tif.SetUpper(upperThreshold)
  tif.SetOutsideValue(outsideValue)
  return tif.Execute(inputImage)

def swt3(inputImage, wavelet="coif1", level=1, start_level=0):
  matrix = sitk.GetArrayFromImage(inputImage)
  matrix = numpy.asarray(matrix)
  data = matrix.copy()
  if data.ndim != 3:
    raise ValueError("Expected 3D data array")

  original_shape = matrix.shape
  adjusted_shape = tuple([dim+1 if dim % 2 != 0 else dim for dim in original_shape])
  data = numpy.resize(data, adjusted_shape)

  if not isinstance(wavelet, pywt.Wavelet):
    wavelet = pywt.Wavelet(wavelet)

  for i in range(0, start_level):
    H, L = decompose_i(data, wavelet)
    LH, LL = decompose_j(L, wavelet)
    LLH, LLL = decompose_k(LL, wavelet)

    data = LLL.copy()

  ret = []
  for i in range(start_level, start_level + level):
    H, L = decompose_i(data, wavelet)

    HH, HL = decompose_j(H, wavelet)
    LH, LL = decompose_j(L, wavelet)

    HHH, HHL = decompose_k(HH, wavelet)
    HLH, HLL = decompose_k(HL, wavelet)
    LHH, LHL = decompose_k(LH, wavelet)
    LLH, LLL = decompose_k(LL, wavelet)

    data = LLL.copy()

    dec = {'HHH': HHH,
           'HHL': HHL,
           'HLH': HLH,
           'HLL': HLL,
           'LHH': LHH,
           'LHL': LHL,
           'LLH': LLH}
    for decName, decImage in dec.iteritems():
      decTemp = decImage.copy()
      decTemp= numpy.resize(decTemp, original_shape)
      sitkImage = sitk.GetImageFromArray(decTemp)
      sitkImage.CopyInformation(inputImage)
      dec[decName] = sitkImage

    ret.append(dec)

  data= numpy.resize(data, original_shape)
  approximation = sitk.GetImageFromArray(data)
  approximation.CopyInformation(inputImage)

  return approximation, ret

def decompose_i(data, wavelet):
  #process in i:
  H, L = [], []
  i_arrays = chain.from_iterable(numpy.transpose(data,(0,1,2)))
  for i_array in i_arrays:
    cA, cD = pywt.swt(i_array, wavelet, level=1, start_level=0)[0]
    H.append(cD)
    L.append(cA)
  H = numpy.hstack(H).reshape(data.shape)
  L = numpy.hstack(L).reshape(data.shape)
  return H, L

def decompose_j(data, wavelet):
  #process in j:
  H, L = [], []
  j_arrays = chain.from_iterable(numpy.transpose(data,(0,1,2)))
  for j_array in j_arrays:
    cA, cD = pywt.swt(j_array, wavelet, level=1, start_level=0)[0]
    H.append(cD)
    L.append(cA)
  H = numpy.asarray( [slice.T for slice in numpy.split(numpy.vstack(H), data.shape[0])] )
  L = numpy.asarray( [slice.T for slice in numpy.split(numpy.vstack(L), data.shape[0])] )
  return H, L

def decompose_k(data, wavelet):
  #process in k:
  H, L = [], []
  k_arrays = chain.from_iterable(numpy.transpose(data,(1,2,0)))
  for k_array in k_arrays:
    cA, cD = pywt.swt(k_array, wavelet, level=1, start_level=0)[0]
    H.append(cD)
    L.append(cA)
  H = numpy.dstack(H).reshape(data.shape)
  L = numpy.dstack(L).reshape(data.shape)
  return H, L

def applySquare(inputImage):
  im = sitk.GetArrayFromImage(inputImage)
  im = im.astype('float64')
  coeff = numpy.sqrt(numpy.max(im)) / numpy.max(im)
  im = (coeff * im) ** 2
  im = sitk.GetImageFromArray(im)
  im.CopyInformation(inputImage)
  return im

def applySquareRoot(inputImage):
  im = sitk.GetArrayFromImage(inputImage)
  im = im.astype('float64')
  coeff = numpy.max(im) ** 2 / numpy.max(im)
  im[im > 0] = numpy.sqrt(im[im > 0] * coeff)
  im[im < 0] = - numpy.sqrt(-im[im < 0] * coeff)
  im = sitk.GetImageFromArray(im)
  im.CopyInformation(inputImage)
  return im

def applyLogarithm(inputImage):
  im = sitk.GetArrayFromImage(inputImage)
  im = im.astype('float64')
  im_max = numpy.max(im)
  im[im > 0] = numpy.log(im[im > 0] + 1)
  im[im < 0] = - numpy.log(- (im[im < 0] - 1))
  im = im * (im_max / numpy.max(im))
  im = sitk.GetImageFromArray(im)
  im.CopyInformation(inputImage)
  return im

def applyExponential(inputImage):
  im = sitk.GetArrayFromImage(inputImage)
  im = im.astype('float64')
  coeff = numpy.log(numpy.max(im))/numpy.max(im)
  im = numpy.exp(coeff*im)
  im = sitk.GetImageFromArray(im)
  im.CopyInformation(inputImage)
  return im
<|endoftext|>


<commit_before># Copyright 2019 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This file is autogenerated by copybara, please do not edit.

bazel_version = "bb41243c5f989c0f6429966b15a9b98c3e7b69eb"
bazel_sha256 = "fe6e56e031f50ad99aec5b4fb4870bc7c901f5ca87b1dabf106a81a7bdbadd23"
<commit_msg>Update dependency bazelbuild/bazel to latest version<commit_after># Copyright 2019 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This file is autogenerated by copybara, please do not edit.

bazel_version = "3bca2ad024d5734a21e64df5a3c57890c820da6f"
bazel_sha256 = "77494b402681d03110120a4c8902a507ecb53a504dca157d7a2cbc9ed06d0360"
<|endoftext|>


<commit_before># Copyright 2019 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This file is autogenerated by copybara, please do not edit.

bazel_version = "f259b8abfd575f544635f57f3bb6678d566ef309"
bazel_sha256 = "7e262ca5f5595a74d75953dfdcb75b271c2561a292972da7f3be449a3e8b28f6"
<commit_msg>Update dependency bazelbuild/bazel to latest version<commit_after># Copyright 2019 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This file is autogenerated by copybara, please do not edit.

bazel_version = "5f944e085292b8d906e7a45ff7821a9493138654"
bazel_sha256 = "e74896c21ec00f6f69197dbde884760cb875d737357d4323356ef694a3e53337"
<|endoftext|>


<commit_before>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import io
import os
import enum
import numpy
import struct
import datetime

from . import base
from . import __about__ as metadata
from .utils import b
from .utils import s

try:
    from . import _speedups
except ImportError:  # pragma: no cover
    _speedups = None

class Mode(enum.IntEnum):
    #: Automatically detect whether the output is a TTY, if so, write ASCII
    #: otherwise write BINARY
    AUTOMATIC = 0
    #: Force writing ASCII
    ASCII = 1
    #: Force writing BINARY
    BINARY = 2

# For backwards compatibility, leave the original references
AUTOMATIC = Mode.AUTOMATIC
ASCII = Mode.ASCII
BINARY = Mode.BINARY

#: Amount of bytes to read while using buffered reading
BUFFER_SIZE = 4096
#: The amount of bytes in the header field
HEADER_SIZE = 80
#: The amount of bytes in the count field
COUNT_SIZE = 4
#: The maximum amount of triangles we can read from binary files
MAX_COUNT = 1e8
#: The header format, can be safely monkeypatched. Limited to 80 characters
HEADER_FORMAT = '{package_name} ({version}) {now} {name}'

class BaseStl(base.BaseMesh):

    @classmethod
    def load(cls, fh, mode=AUTOMATIC, speedups=True):
        '''Load Mesh from STL file

        Automatically detects binary versus ascii STL files.

        :param file fh: The file handle to open
        :param int mode: Automatically detect the filetype or force binary
        '''
        header = fh.read(HEADER_SIZE)
        if not header:
            return

        if isinstance(header, str):  # pragma: no branch
            header = b(header)

        name = ''

        if mode is AUTOMATIC:
            try:
                headerstr = header.decode("ascii").lstrip().lower()
                possibly_ascii = True if headerstr[:5] == 'solid' else False
            except UnicodeError:
                # If the header could not be decoded to string, it is
                # guaranteed to not be ASCII. If the entire header is
                # decoded correctly, it might nevertheless still be binary -
                # a lot of programs fill the header of binary STL's with
                # valid ASCII characters only (e.g. NULL characters).
                possibly_ascii = False

            if possibly_ascii:
                try:
                    name, data = cls._load_ascii(
                        fh, header, speedups=speedups)
                except RuntimeError as exception:
                    (recoverable, e) = exception.args
                    # If we didn't read beyond the header the stream is still
                    # readable through the binary reader
                    if recoverable:
                        name, data = cls._load_binary(fh, header,
                                                      check_size=False)
                    else:
                        # Apparently we've read beyond the header. Let's try
                        # seeking :)
                        # Note that this fails when reading from stdin, we
                        # can't recover from that.
                        fh.seek(HEADER_SIZE)

                        # Since we know this is a seekable file now and we're
                        # not 100% certain it's binary, check the size while
                        # reading
                        name, data = cls._load_binary(fh, header,
                                                      check_size=True)
            else:
                name, data = cls._load_binary(fh, header)
        elif mode is ASCII:
            name, data = cls._load_ascii(fh, header, speedups=speedups)
        else:
            name, data = cls._load_binary(fh, header)

        return name, data

    @classmethod
    def _load_binary(cls, fh, header, check_size=False):
        # Read the triangle count
        count_data = fh.read(COUNT_SIZE)
        if len(count_data) != COUNT_SIZE:
            count = 0
        else:
            count, = struct.unpack(s('<i'), b(count_data))
        # raise RuntimeError()
        assert count < MAX_COUNT, ('File too large, got %d triangles which '
                                   'exceeds the maximum of %d') % (
                                       count, MAX_COUNT)

        if check_size:
            try:
                # Check the size of the file
                fh.seek(0, os.SEEK_END)
                raw_size = fh.tell() - HEADER_SIZE - COUNT_SIZE
                expected_count = int(raw_size / cls.dtype.itemsize)
                assert expected_count == count, ('Expected %d vectors but '
                                                 'header indicates %d') % (
                                                     expected_count, count)
                fh.seek(HEADER_SIZE + COUNT_SIZE)
            except IOError:  # pragma: no cover
                pass

        name = header.strip()

        # Read the rest of the binary data
        try:
            return name, numpy.fromfile(fh, dtype=cls.dtype, count=count)
        except io.UnsupportedOperation:
            data = numpy.frombuffer(fh.read(), dtype=cls.dtype, count=count)
            # Copy to make the buffer writable
            return name, data.copy()

    @classmethod
    def _ascii_reader(cls, fh, header):
        if b'\n' in header:
            recoverable = [True]
        else:
            recoverable = [False]
            header += b(fh.read(BUFFER_SIZE))

        lines = b(header).split(b('\n'))

        def get(prefix=''):
            prefix = b(prefix).lower()

            if lines:
                raw_line = lines.pop(0)
            else:
                raise RuntimeError(recoverable[0], 'Unable to find more lines')

            if not lines:
                recoverable[0] = False

                # Read more lines and make sure we prepend any old data
                lines[:] = b(fh.read(BUFFER_SIZE)).split(b('\n'))
                raw_line += lines.pop(0)

            raw_line = raw_line.strip()
            line = raw_line.lower()
            if line == b(''):
                return get(prefix)

            if prefix:
                if line.startswith(prefix):
                    values = line.replace(prefix, b(''), 1).strip().split()
                elif line.startswith(b('endsolid')):
                    # go back to the beginning of new solid part
                    size_unprocessedlines = sum(
                        len(line) + 1 for line in lines) - 1

                    if size_unprocessedlines > 0:
                        position = fh.tell()
                        fh.seek(position - size_unprocessedlines)
                    raise StopIteration()
                else:
                    raise RuntimeError(
                        recoverable[0],
                        '%r should start with %r' % (line, prefix))

                if len(values) == 3:
                    return [float(v) for v in values]
                else:  # pragma: no cover
                    raise RuntimeError(recoverable[0],
                                       'Incorrect value %r' % line)
            else:
                return b(raw_line)

        line = get()
        if not lines:
            raise RuntimeError(recoverable[0],
                               'No lines found, impossible to read')

        # Yield the name
        yield line[5:].strip()

        while True:
            # Read from the header lines first, until that point we can recover
            # and go to the binary option. After that we cannot due to
            # unseekable files such as sys.stdin
            #
            # Numpy doesn't support any non-file types so wrapping with a
            # buffer and/or StringIO does not work.
            try:
                normals = get('facet normal')
                assert get().lower() == b('outer loop')
                v0 = get('vertex')
                v1 = get('vertex')
                v2 = get('vertex')
                assert get().lower() == b('endloop')
                assert get().lower() == b('endfacet')
                attrs = 0
                yield (normals, (v0, v1, v2), attrs)
            except AssertionError as e:  # pragma: no cover
                raise RuntimeError(recoverable[0], e)
            except StopIteration:
                return

    @classmethod
    def _load_ascii(cls, fh, header, speedups=True):
        # Speedups does not support non file-based streams
        try:
            fh.fileno()
        except io.UnsupportedOperation:
            speedups = False
        # The speedups module is covered by travis but it can't be tested in
        # all environments, this makes coverage checks easier
        if _speedups and speedups:  # pragma: no cover
            return _speedups.ascii_read(fh, header)
        else:
            iterator = cls._ascii_reader(fh, header)
            name = next(iterator)
            return name, numpy.fromiter(iterator, dtype=cls.dtype)

    def save(self, filename, fh=None, mode=AUTOMATIC, update_normals=True):
        '''Save the STL to a (binary) file

        If mode is :py:data:`AUTOMATIC` an :py:data:`ASCII` file will be
        written if the output is a TTY and a :py:data:`BINARY` file otherwise.

        :param str filename: The file to load
        :param file fh: The file handle to open
        :param int mode: The mode to write, default is :py:data:`AUTOMATIC`.
        :param bool update_normals: Whether to update the normals
        '''
        assert filename, 'Filename is required for the STL headers'
        if update_normals:
            self.update_normals()

        if mode is AUTOMATIC:
            if fh and os.isatty(fh.fileno()):  # pragma: no cover
                write = self._write_ascii
            else:
                write = self._write_binary
        elif mode is BINARY:
            write = self._write_binary
        elif mode is ASCII:
            write = self._write_ascii
        else:
            raise ValueError('Mode %r is invalid' % mode)

        name = os.path.split(filename)[-1]
        try:
            if fh:
                write(fh, name)
            else:
                with open(filename, 'wb') as fh:
                    write(fh, filename)
        except IOError:  # pragma: no cover
            pass

    def _write_ascii(self, fh, name):
        if _speedups and self.speedups:  # pragma: no cover
            _speedups.ascii_write(fh, b(name), self.data)
        else:
            def p(s, file):
                file.write(b('%s\n' % s))

            p('solid %s' % name, file=fh)

            for row in self.data:
                vectors = row['vectors']
                p('facet normal %f %f %f' % tuple(row['normals']), file=fh)
                p('  outer loop', file=fh)
                p('    vertex %f %f %f' % tuple(vectors[0]), file=fh)
                p('    vertex %f %f %f' % tuple(vectors[1]), file=fh)
                p('    vertex %f %f %f' % tuple(vectors[2]), file=fh)
                p('  endloop', file=fh)
                p('endfacet', file=fh)

            p('endsolid %s' % name, file=fh)

    def get_header(self, name):
        # Format the header
        header = HEADER_FORMAT.format(
            package_name=metadata.__package_name__,
            version=metadata.__version__,
            now=datetime.datetime.now(),
            name=name,
        )

        # Make it exactly 80 characters
        return header[:80].ljust(80, ' ')

    def _write_binary(self, fh, name):
        header = self.get_header(name)
        packed = struct.pack(s('<i'), self.data.size)

        if isinstance(fh, io.TextIOWrapper):  # pragma: no cover
            packed = str(packed)
        else:
            header = b(header)
            packed = b(packed)

        fh.write(header)
        fh.write(packed)
        self.data.tofile(fh)

        if self.data.size:  # pragma: no cover
            assert fh.tell() > 84, (
                'numpy silently refused to write our file. Note that writing '
                'to `StringIO` objects is not supported by `numpy`')

    @classmethod
    def from_file(cls, filename, calculate_normals=True, fh=None,
                  mode=Mode.AUTOMATIC, speedups=True, **kwargs):
        '''Load a mesh from a STL file

        :param str filename: The file to load
        :param bool calculate_normals: Whether to update the normals
        :param file fh: The file handle to open
        :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`

        '''
        if fh:
            name, data = cls.load(
                fh, mode=mode, speedups=speedups)
        else:
            with open(filename, 'rb') as fh:
                name, data = cls.load(
                    fh, mode=mode, speedups=speedups)

        return cls(data, calculate_normals, name=name,
                   speedups=speedups, **kwargs)

    @classmethod
    def from_multi_file(cls, filename, calculate_normals=True, fh=None,
                        mode=Mode.ASCII, speedups=True, **kwargs):
        '''Load multiple meshes from a STL file

        Note: mode is hardcoded to ascii since binary stl files do not support
        the multi format

        :param str filename: The file to load
        :param bool calculate_normals: Whether to update the normals
        :param file fh: The file handle to open
        :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
        '''
        if fh:
            close = False
        else:
            fh = open(filename, 'rb')
            close = True

        try:
            raw_data = cls.load(fh, mode=mode, speedups=speedups)
            while raw_data:
                name, data = raw_data
                yield cls(data, calculate_normals, name=name,
                          speedups=speedups, **kwargs)
                raw_data = cls.load(fh, mode=ASCII,
                                    speedups=speedups)

        finally:
            if close:
                fh.close()

    @classmethod
    def from_files(cls, filenames, calculate_normals=True, mode=Mode.AUTOMATIC,
                   speedups=True, **kwargs):
        '''Load multiple meshes from a STL file

        Note: mode is hardcoded to ascii since binary stl files do not support
        the multi format

        :param list(str) filenames: The files to load
        :param bool calculate_normals: Whether to update the normals
        :param file fh: The file handle to open
        :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
        '''
        meshes = []
        for filename in filenames:
            meshes.append(cls.from_file(
                filename,
                calculate_normals=calculate_normals,
                mode=mode,
                speedups=speedups,
                **kwargs))

        data = numpy.concatenate([mesh.data for mesh in meshes])
        return cls(data, calculate_normals=calculate_normals, **kwargs)

StlMesh = BaseStl.from_file
<commit_msg>simplified possibly ascii test<commit_after>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import io
import os
import enum
import numpy
import struct
import datetime

from . import base
from . import __about__ as metadata
from .utils import b
from .utils import s

try:
    from . import _speedups
except ImportError:  # pragma: no cover
    _speedups = None

class Mode(enum.IntEnum):
    #: Automatically detect whether the output is a TTY, if so, write ASCII
    #: otherwise write BINARY
    AUTOMATIC = 0
    #: Force writing ASCII
    ASCII = 1
    #: Force writing BINARY
    BINARY = 2

# For backwards compatibility, leave the original references
AUTOMATIC = Mode.AUTOMATIC
ASCII = Mode.ASCII
BINARY = Mode.BINARY

#: Amount of bytes to read while using buffered reading
BUFFER_SIZE = 4096
#: The amount of bytes in the header field
HEADER_SIZE = 80
#: The amount of bytes in the count field
COUNT_SIZE = 4
#: The maximum amount of triangles we can read from binary files
MAX_COUNT = 1e8
#: The header format, can be safely monkeypatched. Limited to 80 characters
HEADER_FORMAT = '{package_name} ({version}) {now} {name}'

class BaseStl(base.BaseMesh):

    @classmethod
    def load(cls, fh, mode=AUTOMATIC, speedups=True):
        '''Load Mesh from STL file

        Automatically detects binary versus ascii STL files.

        :param file fh: The file handle to open
        :param int mode: Automatically detect the filetype or force binary
        '''
        header = fh.read(HEADER_SIZE)
        if not header:
            return

        if isinstance(header, str):  # pragma: no branch
            header = b(header)

        name = ''

        if mode is AUTOMATIC:
            try:
                possibly_ascii = header.lstrip().lower().startswith(b'solid')
            except UnicodeError:
                # If the header could not be decoded to string, it is
                # guaranteed to not be ASCII. If the entire header is
                # decoded correctly, it might nevertheless still be binary -
                # a lot of programs fill the header of binary STL's with
                # valid ASCII characters only (e.g. NULL characters).
                possibly_ascii = False

            if possibly_ascii:
                try:
                    name, data = cls._load_ascii(
                        fh, header, speedups=speedups)
                except RuntimeError as exception:
                    (recoverable, e) = exception.args
                    # If we didn't read beyond the header the stream is still
                    # readable through the binary reader
                    if recoverable:
                        name, data = cls._load_binary(fh, header,
                                                      check_size=False)
                    else:
                        # Apparently we've read beyond the header. Let's try
                        # seeking :)
                        # Note that this fails when reading from stdin, we
                        # can't recover from that.
                        fh.seek(HEADER_SIZE)

                        # Since we know this is a seekable file now and we're
                        # not 100% certain it's binary, check the size while
                        # reading
                        name, data = cls._load_binary(fh, header,
                                                      check_size=True)
            else:
                name, data = cls._load_binary(fh, header)
        elif mode is ASCII:
            name, data = cls._load_ascii(fh, header, speedups=speedups)
        else:
            name, data = cls._load_binary(fh, header)

        return name, data

    @classmethod
    def _load_binary(cls, fh, header, check_size=False):
        # Read the triangle count
        count_data = fh.read(COUNT_SIZE)
        if len(count_data) != COUNT_SIZE:
            count = 0
        else:
            count, = struct.unpack(s('<i'), b(count_data))
        # raise RuntimeError()
        assert count < MAX_COUNT, ('File too large, got %d triangles which '
                                   'exceeds the maximum of %d') % (
                                       count, MAX_COUNT)

        if check_size:
            try:
                # Check the size of the file
                fh.seek(0, os.SEEK_END)
                raw_size = fh.tell() - HEADER_SIZE - COUNT_SIZE
                expected_count = int(raw_size / cls.dtype.itemsize)
                assert expected_count == count, ('Expected %d vectors but '
                                                 'header indicates %d') % (
                                                     expected_count, count)
                fh.seek(HEADER_SIZE + COUNT_SIZE)
            except IOError:  # pragma: no cover
                pass

        name = header.strip()

        # Read the rest of the binary data
        try:
            return name, numpy.fromfile(fh, dtype=cls.dtype, count=count)
        except io.UnsupportedOperation:
            data = numpy.frombuffer(fh.read(), dtype=cls.dtype, count=count)
            # Copy to make the buffer writable
            return name, data.copy()

    @classmethod
    def _ascii_reader(cls, fh, header):
        if b'\n' in header:
            recoverable = [True]
        else:
            recoverable = [False]
            header += b(fh.read(BUFFER_SIZE))

        lines = b(header).split(b('\n'))

        def get(prefix=''):
            prefix = b(prefix).lower()

            if lines:
                raw_line = lines.pop(0)
            else:
                raise RuntimeError(recoverable[0], 'Unable to find more lines')

            if not lines:
                recoverable[0] = False

                # Read more lines and make sure we prepend any old data
                lines[:] = b(fh.read(BUFFER_SIZE)).split(b('\n'))
                raw_line += lines.pop(0)

            raw_line = raw_line.strip()
            line = raw_line.lower()
            if line == b(''):
                return get(prefix)

            if prefix:
                if line.startswith(prefix):
                    values = line.replace(prefix, b(''), 1).strip().split()
                elif line.startswith(b('endsolid')):
                    # go back to the beginning of new solid part
                    size_unprocessedlines = sum(
                        len(line) + 1 for line in lines) - 1

                    if size_unprocessedlines > 0:
                        position = fh.tell()
                        fh.seek(position - size_unprocessedlines)
                    raise StopIteration()
                else:
                    raise RuntimeError(
                        recoverable[0],
                        '%r should start with %r' % (line, prefix))

                if len(values) == 3:
                    return [float(v) for v in values]
                else:  # pragma: no cover
                    raise RuntimeError(recoverable[0],
                                       'Incorrect value %r' % line)
            else:
                return b(raw_line)

        line = get()
        if not lines:
            raise RuntimeError(recoverable[0],
                               'No lines found, impossible to read')

        # Yield the name
        yield line[5:].strip()

        while True:
            # Read from the header lines first, until that point we can recover
            # and go to the binary option. After that we cannot due to
            # unseekable files such as sys.stdin
            #
            # Numpy doesn't support any non-file types so wrapping with a
            # buffer and/or StringIO does not work.
            try:
                normals = get('facet normal')
                assert get().lower() == b('outer loop')
                v0 = get('vertex')
                v1 = get('vertex')
                v2 = get('vertex')
                assert get().lower() == b('endloop')
                assert get().lower() == b('endfacet')
                attrs = 0
                yield (normals, (v0, v1, v2), attrs)
            except AssertionError as e:  # pragma: no cover
                raise RuntimeError(recoverable[0], e)
            except StopIteration:
                return

    @classmethod
    def _load_ascii(cls, fh, header, speedups=True):
        # Speedups does not support non file-based streams
        try:
            fh.fileno()
        except io.UnsupportedOperation:
            speedups = False
        # The speedups module is covered by travis but it can't be tested in
        # all environments, this makes coverage checks easier
        if _speedups and speedups:  # pragma: no cover
            return _speedups.ascii_read(fh, header)
        else:
            iterator = cls._ascii_reader(fh, header)
            name = next(iterator)
            return name, numpy.fromiter(iterator, dtype=cls.dtype)

    def save(self, filename, fh=None, mode=AUTOMATIC, update_normals=True):
        '''Save the STL to a (binary) file

        If mode is :py:data:`AUTOMATIC` an :py:data:`ASCII` file will be
        written if the output is a TTY and a :py:data:`BINARY` file otherwise.

        :param str filename: The file to load
        :param file fh: The file handle to open
        :param int mode: The mode to write, default is :py:data:`AUTOMATIC`.
        :param bool update_normals: Whether to update the normals
        '''
        assert filename, 'Filename is required for the STL headers'
        if update_normals:
            self.update_normals()

        if mode is AUTOMATIC:
            if fh and os.isatty(fh.fileno()):  # pragma: no cover
                write = self._write_ascii
            else:
                write = self._write_binary
        elif mode is BINARY:
            write = self._write_binary
        elif mode is ASCII:
            write = self._write_ascii
        else:
            raise ValueError('Mode %r is invalid' % mode)

        name = os.path.split(filename)[-1]
        try:
            if fh:
                write(fh, name)
            else:
                with open(filename, 'wb') as fh:
                    write(fh, filename)
        except IOError:  # pragma: no cover
            pass

    def _write_ascii(self, fh, name):
        if _speedups and self.speedups:  # pragma: no cover
            _speedups.ascii_write(fh, b(name), self.data)
        else:
            def p(s, file):
                file.write(b('%s\n' % s))

            p('solid %s' % name, file=fh)

            for row in self.data:
                vectors = row['vectors']
                p('facet normal %f %f %f' % tuple(row['normals']), file=fh)
                p('  outer loop', file=fh)
                p('    vertex %f %f %f' % tuple(vectors[0]), file=fh)
                p('    vertex %f %f %f' % tuple(vectors[1]), file=fh)
                p('    vertex %f %f %f' % tuple(vectors[2]), file=fh)
                p('  endloop', file=fh)
                p('endfacet', file=fh)

            p('endsolid %s' % name, file=fh)

    def get_header(self, name):
        # Format the header
        header = HEADER_FORMAT.format(
            package_name=metadata.__package_name__,
            version=metadata.__version__,
            now=datetime.datetime.now(),
            name=name,
        )

        # Make it exactly 80 characters
        return header[:80].ljust(80, ' ')

    def _write_binary(self, fh, name):
        header = self.get_header(name)
        packed = struct.pack(s('<i'), self.data.size)

        if isinstance(fh, io.TextIOWrapper):  # pragma: no cover
            packed = str(packed)
        else:
            header = b(header)
            packed = b(packed)

        fh.write(header)
        fh.write(packed)
        self.data.tofile(fh)

        if self.data.size:  # pragma: no cover
            assert fh.tell() > 84, (
                'numpy silently refused to write our file. Note that writing '
                'to `StringIO` objects is not supported by `numpy`')

    @classmethod
    def from_file(cls, filename, calculate_normals=True, fh=None,
                  mode=Mode.AUTOMATIC, speedups=True, **kwargs):
        '''Load a mesh from a STL file

        :param str filename: The file to load
        :param bool calculate_normals: Whether to update the normals
        :param file fh: The file handle to open
        :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`

        '''
        if fh:
            name, data = cls.load(
                fh, mode=mode, speedups=speedups)
        else:
            with open(filename, 'rb') as fh:
                name, data = cls.load(
                    fh, mode=mode, speedups=speedups)

        return cls(data, calculate_normals, name=name,
                   speedups=speedups, **kwargs)

    @classmethod
    def from_multi_file(cls, filename, calculate_normals=True, fh=None,
                        mode=Mode.ASCII, speedups=True, **kwargs):
        '''Load multiple meshes from a STL file

        Note: mode is hardcoded to ascii since binary stl files do not support
        the multi format

        :param str filename: The file to load
        :param bool calculate_normals: Whether to update the normals
        :param file fh: The file handle to open
        :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
        '''
        if fh:
            close = False
        else:
            fh = open(filename, 'rb')
            close = True

        try:
            raw_data = cls.load(fh, mode=mode, speedups=speedups)
            while raw_data:
                name, data = raw_data
                yield cls(data, calculate_normals, name=name,
                          speedups=speedups, **kwargs)
                raw_data = cls.load(fh, mode=ASCII,
                                    speedups=speedups)

        finally:
            if close:
                fh.close()

    @classmethod
    def from_files(cls, filenames, calculate_normals=True, mode=Mode.AUTOMATIC,
                   speedups=True, **kwargs):
        '''Load multiple meshes from a STL file

        Note: mode is hardcoded to ascii since binary stl files do not support
        the multi format

        :param list(str) filenames: The files to load
        :param bool calculate_normals: Whether to update the normals
        :param file fh: The file handle to open
        :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
        '''
        meshes = []
        for filename in filenames:
            meshes.append(cls.from_file(
                filename,
                calculate_normals=calculate_normals,
                mode=mode,
                speedups=speedups,
                **kwargs))

        data = numpy.concatenate([mesh.data for mesh in meshes])
        return cls(data, calculate_normals=calculate_normals, **kwargs)

StlMesh = BaseStl.from_file
<|endoftext|>