<commit_before>"""Import tasks related to MOSFiT."""
import os
from glob import glob

import dropbox

from astrocats.catalog.utils import pbar
from astrocats.supernovae.supernova import Supernova, SUPERNOVA
from astrocats.catalog.photometry import PHOTOMETRY

def do_mosfit(catalog):
    """Import models produced by MOSFiT from Dropbox."""
    REALIZATION_LIMIT = 10

    task_str = catalog.get_current_task_str()
    try:
        with open('mosfit.key', 'r') as f:
            mosfitkey = f.read().splitlines()[0]
    except Exception:
        catalog.log.warning('MOSFiT key not found, make sure a file named '
                            '`mosfit.key` containing the key is placed the '
                            'astrocats directory.')
        mosfitkey = ''

    dbx = dropbox.Dropbox(mosfitkey)
    files = list(sorted([
        x.name for x in dbx.files_list_folder('').entries
        if not x.name.startswith('.')
    ]))
    fdir = os.path.join(catalog.get_current_task_repo(), 'MOSFiT')
    if not os.path.isdir(fdir):
        os.mkdir(fdir)
    efiles = [x.split('/')[-1] for x in glob(os.path.join(fdir, '*'))]
    old_name = ''
    for fname in pbar(files, desc=task_str):
        if fname in efiles:
            efiles.remove(fname)
        fpath = os.path.join(fdir, fname)
        if not os.path.isfile(fpath):
            md, res = dbx.files_download('/' + fname)
            jtxt = res.content
            with open(fpath, 'wb') as f:
                f.write(jtxt)

        new_entry = Supernova.init_from_file(
            catalog, path=fpath, compare_to_existing=False, try_gzip=True,
            clean=False, merge=False, filter_on={
                'realization': [str(x) for x in range(1, REALIZATION_LIMIT)]})

        name = new_entry[SUPERNOVA.NAME]

        # Only take a number of realizations up to the realization limit.
        new_photo = []
        for photo in new_entry[SUPERNOVA.PHOTOMETRY]:
            real = int(photo.get(PHOTOMETRY.REALIZATION, -1))
            if real < 0 or real >= REALIZATION_LIMIT:
                continue
            new_photo.append(photo)
        new_entry[SUPERNOVA.PHOTOMETRY] = new_photo

        old_entry = None
        if name in catalog.entries:
            if catalog.entries[name]._stub:
                old_entry = Supernova.init_from_file(
                    catalog, name=name, compare_to_existing=False)
            else:
                old_entry = catalog.entries[name]

        if old_entry:
            catalog.copy_entry_to_entry(new_entry, old_entry,
                                        compare_to_existing=False)
            catalog.entries[name] = old_entry
        else:
            catalog.entries[name] = new_entry

        if old_name != name:
            catalog.journal_entries()
        old_name = name
    for fname in efiles:
        os.remove(os.path.join(fdir, fname))

    return
<commit_msg>ENH: update mosfit task<commit_after>"""Import tasks related to MOSFiT."""
import os
from glob import glob

import dropbox

from astrocats.catalog.utils import pbar
from astrocats.supernovae.supernova import Supernova, SUPERNOVA
from astrocats.catalog.photometry import PHOTOMETRY

def do_mosfit(catalog):
    """Import models produced by MOSFiT from Dropbox."""
    REALIZATION_LIMIT = 10

    task_str = catalog.get_current_task_str()
    try:
        with open('mosfit.key', 'r') as f:
            mosfitkey = f.read().splitlines()[0]
    except Exception:
        catalog.log.warning('MOSFiT key not found, make sure a file named '
                            '`mosfit.key` containing the key is placed the '
                            'astrocats directory.')
        mosfitkey = ''

    # Get new data from Dropbox.
    dbx = dropbox.Dropbox(mosfitkey)
    files = list(sorted([
        x.name for x in dbx.files_list_folder('').entries
        if not x.name.startswith('.')
    ]))
    fdir = os.path.join(catalog.get_current_task_repo(), 'MOSFiT')
    if not os.path.isdir(fdir):
        os.mkdir(fdir)
    old_name = ''
    for fname in pbar(files, desc=task_str):
        fpath = os.path.join(fdir, fname)

        if 'GW' in fpath:
            continue

        if not os.path.isfile(fpath):
            md, res = dbx.files_download('/' + fname)
            jtxt = res.content
            with open(fpath, 'wb') as f:
                f.write(jtxt)

    # Load data in models folder.
    efiles = [x.split('/')[-1] for x in glob(
        os.path.join(fdir, '*')) if '.json' in x]
    for fname in efiles:
        fpath = os.path.join(fdir, fname)
        new_entry = Supernova.init_from_file(
            catalog, path=fpath, compare_to_existing=False, try_gzip=True,
            clean=False, merge=False, filter_on={
                'realization': [str(x) for x in range(1, REALIZATION_LIMIT)]})

        name = new_entry[SUPERNOVA.NAME]

        aliases = new_entry.get_aliases()

        if any([x.startswith('GW') for x in aliases]):
            os.remove(fpath)
            continue

        # Only take a number of realizations up to the realization limit.
        new_photo = []
        for photo in new_entry[SUPERNOVA.PHOTOMETRY]:
            real = int(photo.get(PHOTOMETRY.REALIZATION, -1))
            if real < 0 or real >= REALIZATION_LIMIT:
                continue
            new_photo.append(photo)
        new_entry[SUPERNOVA.PHOTOMETRY] = new_photo

        old_entry = None
        if name in catalog.entries:
            if catalog.entries[name]._stub:
                old_entry = Supernova.init_from_file(
                    catalog, name=name, compare_to_existing=False)
            else:
                old_entry = catalog.entries[name]

        if old_entry:
            catalog.copy_entry_to_entry(new_entry, old_entry,
                                        compare_to_existing=False)
            catalog.entries[name] = old_entry
        else:
            catalog.entries[name] = new_entry

        if old_name != name:
            catalog.journal_entries()
        old_name = name

    return
<|endoftext|>


<commit_before>"""Output streaming, processing and formatting.

"""
import json
import xml.dom.minidom
from functools import partial
from itertools import chain

import pygments
from pygments import token, lexer
from pygments.styles import get_style_by_name, STYLE_MAP
from pygments.lexers import get_lexer_for_mimetype, get_lexer_by_name
from pygments.formatters.terminal import TerminalFormatter
from pygments.formatters.terminal256 import Terminal256Formatter
from pygments.util import ClassNotFound

from .compat import is_windows
from .solarized import Solarized256Style
from .models import HTTPRequest, HTTPResponse, Environment
from .input import (OUT_REQ_BODY, OUT_REQ_HEAD,
                    OUT_RESP_HEAD, OUT_RESP_BODY)

# The default number of spaces to indent when pretty printing
DEFAULT_INDENT = 4

# Colors on Windows via colorama don't look that
# great and fruity seems to give the best result there.
AVAILABLE_STYLES = set(STYLE_MAP.keys())
AVAILABLE_STYLES.add('solarized')
DEFAULT_STYLE = 'solarized' if not is_windows else 'fruity'

BINARY_SUPPRESSED_NOTICE = (
    b'\n'
    b'+-----------------------------------------+\n'
    b'| NOTE: binary data not shown in terminal |\n'
    b'+-----------------------------------------+'
)

class BinarySuppressedError(Exception):
    """An error indicating that the body is binary and won't be written,
     e.g., for terminal output)."""

    message = BINARY_SUPPRESSED_NOTICE

###############################################################################
# Output Streams
###############################################################################

def write(stream, outfile, flush):
    """Write the output stream."""
    try:
        # Writing bytes so we use the buffer interface (Python 3).
        buf = outfile.buffer
    except AttributeError:
        buf = outfile

    for chunk in stream:
        buf.write(chunk)
        if flush:
            outfile.flush()

def write_with_colors_win_py3(stream, outfile, flush):
    """Like `write`, but colorized chunks are written as text
    directly to `outfile` to ensure it gets processed by colorama.
    Applies only to Windows with Python 3 and colorized terminal output.

    """
    color = b'\x1b['
    encoding = outfile.encoding
    for chunk in stream:
        if color in chunk:
            outfile.write(chunk.decode(encoding))
        else:
            outfile.buffer.write(chunk)
        if flush:
            outfile.flush()

def build_output_stream(args, env, request, response):
    """Build and return a chain of iterators over the `request`-`response`
    exchange each of which yields `bytes` chunks.

    """

    req_h = OUT_REQ_HEAD in args.output_options
    req_b = OUT_REQ_BODY in args.output_options
    resp_h = OUT_RESP_HEAD in args.output_options
    resp_b = OUT_RESP_BODY in args.output_options
    req = req_h or req_b
    resp = resp_h or resp_b

    output = []
    Stream = get_stream_type(env, args)

    if req:
        output.append(Stream(
            msg=HTTPRequest(request),
            with_headers=req_h,
            with_body=req_b))

    if req_b and resp:
        # Request/Response separator.
        output.append([b'\n\n'])

    if resp:
        output.append(Stream(
            msg=HTTPResponse(response),
            with_headers=resp_h,
            with_body=resp_b))

    if env.stdout_isatty and resp_b:
        # Ensure a blank line after the response body.
        # For terminal output only.
        output.append([b'\n\n'])

    return chain(*output)

def get_stream_type(env, args):
    """Pick the right stream type based on `env` and `args`.
    Wrap it in a partial with the type-specific args so that
    we don't need to think what stream we are dealing with.

    """
    if not env.stdout_isatty and not args.prettify:
        Stream = partial(
            RawStream,
            chunk_size=RawStream.CHUNK_SIZE_BY_LINE
                       if args.stream
                       else RawStream.CHUNK_SIZE
        )
    elif args.prettify:
        Stream = partial(
            PrettyStream if args.stream else BufferedPrettyStream,
            env=env,
            processor=OutputProcessor(
                env=env, groups=args.prettify, pygments_style=args.style),
        )
    else:
        Stream = partial(EncodedStream, env=env)

    return Stream

class BaseStream(object):
    """Base HTTP message output stream class."""

    def __init__(self, msg, with_headers=True, with_body=True,
                 on_body_chunk_downloaded=None):
        """
        :param msg: a :class:`models.HTTPMessage` subclass
        :param with_headers: if `True`, headers will be included
        :param with_body: if `True`, body will be included

        """
        assert with_headers or with_body
        self.msg = msg
        self.with_headers = with_headers
        self.with_body = with_body
        self.on_body_chunk_downloaded = on_body_chunk_downloaded

    def _get_headers(self):
        """Return the headers' bytes."""
        return self.msg.headers.encode('ascii')

    def _iter_body(self):
        """Return an iterator over the message body."""
        raise NotImplementedError()

    def __iter__(self):
        """Return an iterator over `self.msg`."""
        if self.with_headers:
            yield self._get_headers()
            yield b'\r\n\r\n'

        if self.with_body:
            try:
                for chunk in self._iter_body():
                    yield chunk
                    if self.on_body_chunk_downloaded:
                        self.on_body_chunk_downloaded(chunk)
            except BinarySuppressedError as e:
                if self.with_headers:
                    yield b'\n'
                yield e.message

class RawStream(BaseStream):
    """The message is streamed in chunks with no processing."""

    CHUNK_SIZE = 1024 * 100
    CHUNK_SIZE_BY_LINE = 1

    def __init__(self, chunk_size=CHUNK_SIZE, **kwargs):
        super(RawStream, self).__init__(**kwargs)
        self.chunk_size = chunk_size

    def _iter_body(self):
        return self.msg.iter_body(self.chunk_size)

class EncodedStream(BaseStream):
    """Encoded HTTP message stream.

    The message bytes are converted to an encoding suitable for
    `self.env.stdout`. Unicode errors are replaced and binary data
    is suppressed. The body is always streamed by line.

    """
    CHUNK_SIZE = 1

    def __init__(self, env=Environment(), **kwargs):

        super(EncodedStream, self).__init__(**kwargs)

        if env.stdout_isatty:
            # Use the encoding supported by the terminal.
            output_encoding = getattr(env.stdout, 'encoding', None)
        else:
            # Preserve the message encoding.
            output_encoding = self.msg.encoding

        # Default to utf8 when unsure.
        self.output_encoding = output_encoding or 'utf8'

    def _iter_body(self):

        for line, lf in self.msg.iter_lines(self.CHUNK_SIZE):

            if b'\0' in line:
                raise BinarySuppressedError()

            yield line.decode(self.msg.encoding)\
                      .encode(self.output_encoding, 'replace') + lf

class PrettyStream(EncodedStream):
    """In addition to :class:`EncodedStream` behaviour, this stream applies
    content processing.

    Useful for long-lived HTTP responses that stream by lines
    such as the Twitter streaming API.

    """

    CHUNK_SIZE = 1

    def __init__(self, processor, **kwargs):
        super(PrettyStream, self).__init__(**kwargs)
        self.processor = processor

    def _get_headers(self):
        return self.processor.process_headers(
            self.msg.headers).encode(self.output_encoding)

    def _iter_body(self):
        for line, lf in self.msg.iter_lines(self.CHUNK_SIZE):
            if b'\0' in line:
                raise BinarySuppressedError()
            yield self._process_body(line) + lf

    def _process_body(self, chunk):
        return (self.processor
                    .process_body(
                        content=chunk.decode(self.msg.encoding, 'replace'),
                        content_type=self.msg.content_type,
                        encoding=self.msg.encoding)
                    .encode(self.output_encoding, 'replace'))

class BufferedPrettyStream(PrettyStream):
    """The same as :class:`PrettyStream` except that the body is fully
    fetched before it's processed.

    Suitable regular HTTP responses.

    """

    CHUNK_SIZE = 1024 * 10

    def _iter_body(self):

        # Read the whole body before prettifying it,
        # but bail out immediately if the body is binary.
        body = bytearray()
        for chunk in self.msg.iter_body(self.CHUNK_SIZE):
            if b'\0' in chunk:
                raise BinarySuppressedError()
            body.extend(chunk)

        yield self._process_body(body)

###############################################################################
# Processing
###############################################################################

class HTTPLexer(lexer.RegexLexer):
    """Simplified HTTP lexer for Pygments.

    It only operates on headers and provides a stronger contrast between
    their names and values than the original one bundled with Pygments
    (:class:`pygments.lexers.text import HttpLexer`), especially when
    Solarized color scheme is used.

    """
    name = 'HTTP'
    aliases = ['http']
    filenames = ['*.http']
    tokens = {
        'root': [
            # Request-Line
            (r'([A-Z]+)( +)([^ ]+)( +)(HTTP)(/)(\d+\.\d+)',
             lexer.bygroups(
                 token.Name.Function,
                 token.Text,
                 token.Name.Namespace,
                 token.Text,
                 token.Keyword.Reserved,
                 token.Operator,
                 token.Number
             )),
            # Response Status-Line
            (r'(HTTP)(/)(\d+\.\d+)( +)(\d{3})( +)(.+)',
             lexer.bygroups(
                 token.Keyword.Reserved,  # 'HTTP'
                 token.Operator,  # '/'
                 token.Number,  # Version
                 token.Text,
                 token.Number,  # Status code
                 token.Text,
                 token.Name.Exception,  # Reason
             )),
            # Header
            (r'(.*?)( *)(:)( *)(.+)', lexer.bygroups(
                token.Name.Attribute,  # Name
                token.Text,
                token.Operator,  # Colon
                token.Text,
                token.String  # Value
            ))
        ]
    }

class BaseProcessor(object):
    """Base, noop output processor class."""

    enabled = True

    def __init__(self, env=Environment(), **kwargs):
        """
        :param env: an class:`Environment` instance
        :param kwargs: additional keyword argument that some
                       processor might require.

        """
        self.env = env
        self.kwargs = kwargs

    def process_headers(self, headers):
        """Return processed `headers`

        :param headers: The headers as text.

        """
        return headers

    def process_body(self, content, content_type, subtype, encoding):
        """Return processed `content`.

        :param content: The body content as text
        :param content_type: Full content type, e.g., 'application/atom+xml'.
        :param subtype: E.g. 'xml'.
        :param encoding: The original content encoding.

        """
        return content

class JSONProcessor(BaseProcessor):
    """JSON body processor."""

    def process_body(self, content, content_type, subtype, encoding):
        if subtype == 'json':
            try:
                # Indent the JSON data, sort keys by name, and
                # avoid unicode escapes to improve readability.
                content = json.dumps(json.loads(content),
                                     sort_keys=True,
                                     ensure_ascii=False,
                                     indent=DEFAULT_INDENT)
            except ValueError:
                # Invalid JSON but we don't care.
                pass
        return content

class XMLProcessor(BaseProcessor):
    """XML body processor."""
    # TODO: tests

    def process_body(self, content, content_type, subtype, encoding):
        if subtype == 'xml':
            try:
                # Pretty print the XML; pre-process content into clean string
                raw_string = ''.join(
                    (x.strip() for x in content.encode(encoding).split('\n')))
                doc = xml.dom.minidom.parseString(raw_string)
                content = doc.toprettyxml(indent=' ' * DEFAULT_INDENT)
            except xml.parsers.expat.ExpatError:
                # Ignore invalid XML errors (skips attempting to pretty print)
                pass
        return content

class PygmentsProcessor(BaseProcessor):
    """A processor that applies syntax-highlighting using Pygments
    to the headers, and to the body as well if its content type is recognized.

    """
    def __init__(self, *args, **kwargs):
        super(PygmentsProcessor, self).__init__(*args, **kwargs)

        # Cache that speeds up when we process streamed body by line.
        self.lexers_by_type = {}

        if not self.env.colors:
            self.enabled = False
            return

        try:
            style = get_style_by_name(
                self.kwargs.get('pygments_style', DEFAULT_STYLE))
        except ClassNotFound:
            style = Solarized256Style

        if self.env.is_windows or self.env.colors == 256:
            fmt_class = Terminal256Formatter
        else:
            fmt_class = TerminalFormatter
        self.formatter = fmt_class(style=style)

    def process_headers(self, headers):
        return pygments.highlight(
            headers, HTTPLexer(), self.formatter).strip()

    def process_body(self, content, content_type, subtype, encoding):
        try:
            lexer = self.lexers_by_type.get(content_type)
            if not lexer:
                try:
                    lexer = get_lexer_for_mimetype(content_type)
                except ClassNotFound:
                    lexer = get_lexer_by_name(subtype)
                self.lexers_by_type[content_type] = lexer
        except ClassNotFound:
            pass
        else:
            content = pygments.highlight(content, lexer, self.formatter)
        return content.strip()

class HeadersProcessor(BaseProcessor):
    """Sorts headers by name retaining relative order of multiple headers
    with the same name.

    """
    def process_headers(self, headers):
        lines = headers.splitlines()
        headers = sorted(lines[1:], key=lambda h: h.split(':')[0])
        return '\r\n'.join(lines[:1] + headers)

class OutputProcessor(object):
    """A delegate class that invokes the actual processors."""

    installed_processors = {
        'format': [
            HeadersProcessor,
            JSONProcessor,
            XMLProcessor
        ],
        'colors': [
            PygmentsProcessor
        ]
    }

    def __init__(self, groups, env=Environment(), **kwargs):
        """
        :param env: a :class:`models.Environment` instance
        :param groups: the groups of processors to be applied
        :param kwargs: additional keyword arguments for processors

        """
        self.processors = []
        for group in groups:
            for cls in self.installed_processors[group]:
                processor = cls(env, **kwargs)
                if processor.enabled:
                    self.processors.append(processor)

    def process_headers(self, headers):
        for processor in self.processors:
            headers = processor.process_headers(headers)
        return headers

    def process_body(self, content, content_type, encoding):
        # e.g., 'application/atom+xml'
        content_type = content_type.split(';')[0]
        # e.g., 'xml'
        subtype = content_type.split('/')[-1].split('+')[-1]

        for processor in self.processors:
            content = processor.process_body(
                content,
                content_type,
                subtype,
                encoding
            )

        return content
<commit_msg>replace XML processor with ElementTree with custom indentation<commit_after>"""Output streaming, processing and formatting.

"""
import json
from xml.etree import ElementTree
from functools import partial
from itertools import chain

import pygments
from pygments import token, lexer
from pygments.styles import get_style_by_name, STYLE_MAP
from pygments.lexers import get_lexer_for_mimetype, get_lexer_by_name
from pygments.formatters.terminal import TerminalFormatter
from pygments.formatters.terminal256 import Terminal256Formatter
from pygments.util import ClassNotFound

from .compat import is_windows
from .solarized import Solarized256Style
from .models import HTTPRequest, HTTPResponse, Environment
from .input import (OUT_REQ_BODY, OUT_REQ_HEAD,
                    OUT_RESP_HEAD, OUT_RESP_BODY)

# The default number of spaces to indent when pretty printing
DEFAULT_INDENT = 4

# Colors on Windows via colorama don't look that
# great and fruity seems to give the best result there.
AVAILABLE_STYLES = set(STYLE_MAP.keys())
AVAILABLE_STYLES.add('solarized')
DEFAULT_STYLE = 'solarized' if not is_windows else 'fruity'

BINARY_SUPPRESSED_NOTICE = (
    b'\n'
    b'+-----------------------------------------+\n'
    b'| NOTE: binary data not shown in terminal |\n'
    b'+-----------------------------------------+'
)

class BinarySuppressedError(Exception):
    """An error indicating that the body is binary and won't be written,
     e.g., for terminal output)."""

    message = BINARY_SUPPRESSED_NOTICE

###############################################################################
# Output Streams
###############################################################################

def write(stream, outfile, flush):
    """Write the output stream."""
    try:
        # Writing bytes so we use the buffer interface (Python 3).
        buf = outfile.buffer
    except AttributeError:
        buf = outfile

    for chunk in stream:
        buf.write(chunk)
        if flush:
            outfile.flush()

def write_with_colors_win_py3(stream, outfile, flush):
    """Like `write`, but colorized chunks are written as text
    directly to `outfile` to ensure it gets processed by colorama.
    Applies only to Windows with Python 3 and colorized terminal output.

    """
    color = b'\x1b['
    encoding = outfile.encoding
    for chunk in stream:
        if color in chunk:
            outfile.write(chunk.decode(encoding))
        else:
            outfile.buffer.write(chunk)
        if flush:
            outfile.flush()

def build_output_stream(args, env, request, response):
    """Build and return a chain of iterators over the `request`-`response`
    exchange each of which yields `bytes` chunks.

    """

    req_h = OUT_REQ_HEAD in args.output_options
    req_b = OUT_REQ_BODY in args.output_options
    resp_h = OUT_RESP_HEAD in args.output_options
    resp_b = OUT_RESP_BODY in args.output_options
    req = req_h or req_b
    resp = resp_h or resp_b

    output = []
    Stream = get_stream_type(env, args)

    if req:
        output.append(Stream(
            msg=HTTPRequest(request),
            with_headers=req_h,
            with_body=req_b))

    if req_b and resp:
        # Request/Response separator.
        output.append([b'\n\n'])

    if resp:
        output.append(Stream(
            msg=HTTPResponse(response),
            with_headers=resp_h,
            with_body=resp_b))

    if env.stdout_isatty and resp_b:
        # Ensure a blank line after the response body.
        # For terminal output only.
        output.append([b'\n\n'])

    return chain(*output)

def get_stream_type(env, args):
    """Pick the right stream type based on `env` and `args`.
    Wrap it in a partial with the type-specific args so that
    we don't need to think what stream we are dealing with.

    """
    if not env.stdout_isatty and not args.prettify:
        Stream = partial(
            RawStream,
            chunk_size=RawStream.CHUNK_SIZE_BY_LINE
                       if args.stream
                       else RawStream.CHUNK_SIZE
        )
    elif args.prettify:
        Stream = partial(
            PrettyStream if args.stream else BufferedPrettyStream,
            env=env,
            processor=OutputProcessor(
                env=env, groups=args.prettify, pygments_style=args.style),
        )
    else:
        Stream = partial(EncodedStream, env=env)

    return Stream

class BaseStream(object):
    """Base HTTP message output stream class."""

    def __init__(self, msg, with_headers=True, with_body=True,
                 on_body_chunk_downloaded=None):
        """
        :param msg: a :class:`models.HTTPMessage` subclass
        :param with_headers: if `True`, headers will be included
        :param with_body: if `True`, body will be included

        """
        assert with_headers or with_body
        self.msg = msg
        self.with_headers = with_headers
        self.with_body = with_body
        self.on_body_chunk_downloaded = on_body_chunk_downloaded

    def _get_headers(self):
        """Return the headers' bytes."""
        return self.msg.headers.encode('ascii')

    def _iter_body(self):
        """Return an iterator over the message body."""
        raise NotImplementedError()

    def __iter__(self):
        """Return an iterator over `self.msg`."""
        if self.with_headers:
            yield self._get_headers()
            yield b'\r\n\r\n'

        if self.with_body:
            try:
                for chunk in self._iter_body():
                    yield chunk
                    if self.on_body_chunk_downloaded:
                        self.on_body_chunk_downloaded(chunk)
            except BinarySuppressedError as e:
                if self.with_headers:
                    yield b'\n'
                yield e.message

class RawStream(BaseStream):
    """The message is streamed in chunks with no processing."""

    CHUNK_SIZE = 1024 * 100
    CHUNK_SIZE_BY_LINE = 1

    def __init__(self, chunk_size=CHUNK_SIZE, **kwargs):
        super(RawStream, self).__init__(**kwargs)
        self.chunk_size = chunk_size

    def _iter_body(self):
        return self.msg.iter_body(self.chunk_size)

class EncodedStream(BaseStream):
    """Encoded HTTP message stream.

    The message bytes are converted to an encoding suitable for
    `self.env.stdout`. Unicode errors are replaced and binary data
    is suppressed. The body is always streamed by line.

    """
    CHUNK_SIZE = 1

    def __init__(self, env=Environment(), **kwargs):

        super(EncodedStream, self).__init__(**kwargs)

        if env.stdout_isatty:
            # Use the encoding supported by the terminal.
            output_encoding = getattr(env.stdout, 'encoding', None)
        else:
            # Preserve the message encoding.
            output_encoding = self.msg.encoding

        # Default to utf8 when unsure.
        self.output_encoding = output_encoding or 'utf8'

    def _iter_body(self):

        for line, lf in self.msg.iter_lines(self.CHUNK_SIZE):

            if b'\0' in line:
                raise BinarySuppressedError()

            yield line.decode(self.msg.encoding)\
                      .encode(self.output_encoding, 'replace') + lf

class PrettyStream(EncodedStream):
    """In addition to :class:`EncodedStream` behaviour, this stream applies
    content processing.

    Useful for long-lived HTTP responses that stream by lines
    such as the Twitter streaming API.

    """

    CHUNK_SIZE = 1

    def __init__(self, processor, **kwargs):
        super(PrettyStream, self).__init__(**kwargs)
        self.processor = processor

    def _get_headers(self):
        return self.processor.process_headers(
            self.msg.headers).encode(self.output_encoding)

    def _iter_body(self):
        for line, lf in self.msg.iter_lines(self.CHUNK_SIZE):
            if b'\0' in line:
                raise BinarySuppressedError()
            yield self._process_body(line) + lf

    def _process_body(self, chunk):
        return (self.processor
                    .process_body(
                        content=chunk.decode(self.msg.encoding, 'replace'),
                        content_type=self.msg.content_type,
                        encoding=self.msg.encoding)
                    .encode(self.output_encoding, 'replace'))

class BufferedPrettyStream(PrettyStream):
    """The same as :class:`PrettyStream` except that the body is fully
    fetched before it's processed.

    Suitable regular HTTP responses.

    """

    CHUNK_SIZE = 1024 * 10

    def _iter_body(self):

        # Read the whole body before prettifying it,
        # but bail out immediately if the body is binary.
        body = bytearray()
        for chunk in self.msg.iter_body(self.CHUNK_SIZE):
            if b'\0' in chunk:
                raise BinarySuppressedError()
            body.extend(chunk)

        yield self._process_body(body)

###############################################################################
# Processing
###############################################################################

class HTTPLexer(lexer.RegexLexer):
    """Simplified HTTP lexer for Pygments.

    It only operates on headers and provides a stronger contrast between
    their names and values than the original one bundled with Pygments
    (:class:`pygments.lexers.text import HttpLexer`), especially when
    Solarized color scheme is used.

    """
    name = 'HTTP'
    aliases = ['http']
    filenames = ['*.http']
    tokens = {
        'root': [
            # Request-Line
            (r'([A-Z]+)( +)([^ ]+)( +)(HTTP)(/)(\d+\.\d+)',
             lexer.bygroups(
                 token.Name.Function,
                 token.Text,
                 token.Name.Namespace,
                 token.Text,
                 token.Keyword.Reserved,
                 token.Operator,
                 token.Number
             )),
            # Response Status-Line
            (r'(HTTP)(/)(\d+\.\d+)( +)(\d{3})( +)(.+)',
             lexer.bygroups(
                 token.Keyword.Reserved,  # 'HTTP'
                 token.Operator,  # '/'
                 token.Number,  # Version
                 token.Text,
                 token.Number,  # Status code
                 token.Text,
                 token.Name.Exception,  # Reason
             )),
            # Header
            (r'(.*?)( *)(:)( *)(.+)', lexer.bygroups(
                token.Name.Attribute,  # Name
                token.Text,
                token.Operator,  # Colon
                token.Text,
                token.String  # Value
            ))
        ]
    }

class BaseProcessor(object):
    """Base, noop output processor class."""

    enabled = True

    def __init__(self, env=Environment(), **kwargs):
        """
        :param env: an class:`Environment` instance
        :param kwargs: additional keyword argument that some
                       processor might require.

        """
        self.env = env
        self.kwargs = kwargs

    def process_headers(self, headers):
        """Return processed `headers`

        :param headers: The headers as text.

        """
        return headers

    def process_body(self, content, content_type, subtype, encoding):
        """Return processed `content`.

        :param content: The body content as text
        :param content_type: Full content type, e.g., 'application/atom+xml'.
        :param subtype: E.g. 'xml'.
        :param encoding: The original content encoding.

        """
        return content

class JSONProcessor(BaseProcessor):
    """JSON body processor."""

    def process_body(self, content, content_type, subtype, encoding):
        if subtype == 'json':
            try:
                # Indent the JSON data, sort keys by name, and
                # avoid unicode escapes to improve readability.
                content = json.dumps(json.loads(content),
                                     sort_keys=True,
                                     ensure_ascii=False,
                                     indent=DEFAULT_INDENT)
            except ValueError:
                # Invalid JSON but we don't care.
                pass
        return content

class XMLProcessor(BaseProcessor):
    """XML body processor."""
    # TODO: tests

    # in-place prettyprint formatter
    # c.f. http://effbot.org/zone/element-lib.htm#prettyprint
    @staticmethod
    def indent(elem, indent_text=' ' * DEFAULT_INDENT):
        def _indent(elem, level=0):
            i = "\n" + level * indent_text
            if len(elem):
                if not elem.text or not elem.text.strip():
                    elem.text = i + indent_text
                if not elem.tail or not elem.tail.strip():
                    elem.tail = i
                for elem in elem:
                    _indent(elem, level + 1)
                if not elem.tail or not elem.tail.strip():
                    elem.tail = i
            else:
                if level and (not elem.tail or not elem.tail.strip()):
                    elem.tail = i
        return _indent(elem)

    def process_body(self, content, content_type, subtype, encoding):
        if subtype == 'xml':
            try:
                root = ElementTree.fromstring(content.encode(encoding))
                self.indent(root)
                content = ElementTree.tostring(root)
            except ElementTree.ParseError:
                # Ignore invalid XML errors (skips attempting to pretty print)
                pass
        return content

class PygmentsProcessor(BaseProcessor):
    """A processor that applies syntax-highlighting using Pygments
    to the headers, and to the body as well if its content type is recognized.

    """
    def __init__(self, *args, **kwargs):
        super(PygmentsProcessor, self).__init__(*args, **kwargs)

        # Cache that speeds up when we process streamed body by line.
        self.lexers_by_type = {}

        if not self.env.colors:
            self.enabled = False
            return

        try:
            style = get_style_by_name(
                self.kwargs.get('pygments_style', DEFAULT_STYLE))
        except ClassNotFound:
            style = Solarized256Style

        if self.env.is_windows or self.env.colors == 256:
            fmt_class = Terminal256Formatter
        else:
            fmt_class = TerminalFormatter
        self.formatter = fmt_class(style=style)

    def process_headers(self, headers):
        return pygments.highlight(
            headers, HTTPLexer(), self.formatter).strip()

    def process_body(self, content, content_type, subtype, encoding):
        try:
            lexer = self.lexers_by_type.get(content_type)
            if not lexer:
                try:
                    lexer = get_lexer_for_mimetype(content_type)
                except ClassNotFound:
                    lexer = get_lexer_by_name(subtype)
                self.lexers_by_type[content_type] = lexer
        except ClassNotFound:
            pass
        else:
            content = pygments.highlight(content, lexer, self.formatter)
        return content.strip()

class HeadersProcessor(BaseProcessor):
    """Sorts headers by name retaining relative order of multiple headers
    with the same name.

    """
    def process_headers(self, headers):
        lines = headers.splitlines()
        headers = sorted(lines[1:], key=lambda h: h.split(':')[0])
        return '\r\n'.join(lines[:1] + headers)

class OutputProcessor(object):
    """A delegate class that invokes the actual processors."""

    installed_processors = {
        'format': [
            HeadersProcessor,
            JSONProcessor,
            XMLProcessor
        ],
        'colors': [
            PygmentsProcessor
        ]
    }

    def __init__(self, groups, env=Environment(), **kwargs):
        """
        :param env: a :class:`models.Environment` instance
        :param groups: the groups of processors to be applied
        :param kwargs: additional keyword arguments for processors

        """
        self.processors = []
        for group in groups:
            for cls in self.installed_processors[group]:
                processor = cls(env, **kwargs)
                if processor.enabled:
                    self.processors.append(processor)

    def process_headers(self, headers):
        for processor in self.processors:
            headers = processor.process_headers(headers)
        return headers

    def process_body(self, content, content_type, encoding):
        # e.g., 'application/atom+xml'
        content_type = content_type.split(';')[0]
        # e.g., 'xml'
        subtype = content_type.split('/')[-1].split('+')[-1]

        for processor in self.processors:
            content = processor.process_body(
                content,
                content_type,
                subtype,
                encoding
            )

        return content
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import base64

from lxml.html import fromstring, tostring, HTMLParser

import StringIO
from PIL import Image as ImageLib

from urlparse import urlparse

from scrapy.http import Request
from scrapy.pipelines.media import MediaPipeline

from ..log import logger

class ImgException(Exception):
    pass

class Image():
    IMAGE_MAX_WIDTH = 800

    def __init__(self, data):
        self._image = ImageLib.open(StringIO.StringIO(data))

    @property
    def size(self):
        return self._image.size

    @property
    def type(self):
        return self._image.format

    def optimize(self, q=75):
        image = self._image
        w, h = self._image.size
        if w > self.IMAGE_MAX_WIDTH:
            h = int(float(h)/w*self.IMAGE_MAX_WIDTH)
            w = self.IMAGE_MAX_WIDTH
            image = self._image.resize((w, h), ImageLib.ANTIALIAS)
        buf = StringIO.StringIO()
        image.save(buf, format=self._image.format, quality=q)
        return buf.getvalue()

class ImagesDlownloadPipeline(MediaPipeline):
    MEDIA_NAME = 'image'
    DEFAULT_ITEM_CONTENT_FIELD = 'content'
    IMAGE_MAX_SIZE = 1024*256

    @classmethod
    def from_settings(cls, settings):
        cls.ITEM_CONTENT_FIELD = settings.get('ITEM_CONTENT_FIELD',
                                              cls.DEFAULT_ITEM_CONTENT_FIELD)
        return cls()

    def get_media_requests(self, item, info):
        try:
            doc = fromstring(item[self.ITEM_CONTENT_FIELD],
                             parser=HTMLParser(encoding=item['encoding']))
        except TypeError:
            return []
        try:
            attr = self.spiderinfo.spider.image_url_attr
        except AttributeError:
            attr = 'src'
        urls = []
        for e in doc.xpath('//img'):
            if attr in e.attrib:
                url = e.get(attr).strip()
                if url.startswith('/'):
                    url = urlparse.urljoin(item['link'].strip(), url)
                urls.append((url, e))

        item._doc = doc
        return [Request(url, meta={'img': e})
                for (url, e) in urls if not url.startswith('data')]

    def media_failed(self, failure, request, info):
        logger.error('spider {} image download failed : {}'.format(
            self.spiderinfo.spider.name, request.url))
        try:
            attr = self.spiderinfo.spider.image_url_attr
            img = request.meta['img']
            src = img.get(attr)
            img.set('src', src)
        except AttributeError:
            pass

    def media_downloaded(self, response, request, info):
        if response.status != 200:
            raise ImgException(
                'image status is {}'.format(response.status))
        if not response.body:
            raise ImgException('image size is 0')
        img = response.meta['img']
        src = response.url
        data = response.body
        imglen = len(data)
        img.set('src', src)
        try:
            image = Image(data)
            w, _ = image.size
            if w < 400:
                style = img.get('style').strip()
                if style is None:
                    style = ''
                style += 'float: right;'
                img.set('style', style)
            if imglen > self.IMAGE_MAX_SIZE:
                data = image.optimize()
            imgtype = image.type
        except:
            try:
                imgtype = response.headers['Content-Type'].split('/')[-1]
            except KeyError:
                logger.warning(
                    'spider {} not found Content-Type: {}'.format(
                        self.spiderinfo.spider.name, src))
                return
        img.set('source', src)
        data = base64.b64encode(data)
        img.set('src', 'data:image/{};base64,{}'.format(imgtype, data))

    def item_completed(self, results, item, info):
        item[self.ITEM_CONTENT_FIELD] = tostring(item._doc, pretty_print=True)
        del item._doc
        return item
<commit_msg>urlparse import for 3<commit_after># -*- coding: utf-8 -*-

import base64

from lxml.html import fromstring, tostring, HTMLParser

import StringIO
from PIL import Image as ImageLib

try:
    from urllib.parse import urlparse
except ImportError:
    from urlparse import urlparse

from scrapy.http import Request
from scrapy.pipelines.media import MediaPipeline

from ..log import logger

class ImgException(Exception):
    pass

class Image():
    IMAGE_MAX_WIDTH = 800

    def __init__(self, data):
        self._image = ImageLib.open(StringIO.StringIO(data))

    @property
    def size(self):
        return self._image.size

    @property
    def type(self):
        return self._image.format

    def optimize(self, q=75):
        image = self._image
        w, h = self._image.size
        if w > self.IMAGE_MAX_WIDTH:
            h = int(float(h)/w*self.IMAGE_MAX_WIDTH)
            w = self.IMAGE_MAX_WIDTH
            image = self._image.resize((w, h), ImageLib.ANTIALIAS)
        buf = StringIO.StringIO()
        image.save(buf, format=self._image.format, quality=q)
        return buf.getvalue()

class ImagesDlownloadPipeline(MediaPipeline):
    MEDIA_NAME = 'image'
    DEFAULT_ITEM_CONTENT_FIELD = 'content'
    IMAGE_MAX_SIZE = 1024*256

    @classmethod
    def from_settings(cls, settings):
        cls.ITEM_CONTENT_FIELD = settings.get('ITEM_CONTENT_FIELD',
                                              cls.DEFAULT_ITEM_CONTENT_FIELD)
        return cls()

    def get_media_requests(self, item, info):
        try:
            doc = fromstring(item[self.ITEM_CONTENT_FIELD],
                             parser=HTMLParser(encoding=item['encoding']))
        except TypeError:
            return []
        try:
            attr = self.spiderinfo.spider.image_url_attr
        except AttributeError:
            attr = 'src'
        urls = []
        for e in doc.xpath('//img'):
            if attr in e.attrib:
                url = e.get(attr).strip()
                if url.startswith('/'):
                    url = urlparse.urljoin(item['link'].strip(), url)
                urls.append((url, e))

        item._doc = doc
        return [Request(url, meta={'img': e})
                for (url, e) in urls if not url.startswith('data')]

    def media_failed(self, failure, request, info):
        logger.error('spider {} image download failed : {}'.format(
            self.spiderinfo.spider.name, request.url))
        try:
            attr = self.spiderinfo.spider.image_url_attr
            img = request.meta['img']
            src = img.get(attr)
            img.set('src', src)
        except AttributeError:
            pass

    def media_downloaded(self, response, request, info):
        if response.status != 200:
            raise ImgException(
                'image status is {}'.format(response.status))
        if not response.body:
            raise ImgException('image size is 0')
        img = response.meta['img']
        src = response.url
        data = response.body
        imglen = len(data)
        img.set('src', src)
        try:
            image = Image(data)
            w, _ = image.size
            if w < 400:
                style = img.get('style').strip()
                if style is None:
                    style = ''
                style += 'float: right;'
                img.set('style', style)
            if imglen > self.IMAGE_MAX_SIZE:
                data = image.optimize()
            imgtype = image.type
        except:
            try:
                imgtype = response.headers['Content-Type'].split('/')[-1]
            except KeyError:
                logger.warning(
                    'spider {} not found Content-Type: {}'.format(
                        self.spiderinfo.spider.name, src))
                return
        img.set('source', src)
        data = base64.b64encode(data)
        img.set('src', 'data:image/{};base64,{}'.format(imgtype, data))

    def item_completed(self, results, item, info):
        item[self.ITEM_CONTENT_FIELD] = tostring(item._doc, pretty_print=True)
        del item._doc
        return item
<|endoftext|>


<commit_before>import os
import errno
from jinja2 import Template

_fsm_header_template = Template("""\
#ifndef {{ header.basename|upper }}_H
#define {{ header.basename|upper }}_H

/* ***************************
 * Typedefs and Declarations *
 * ***************************/

{% for code_snippet in header.code -%}
{{ code_snippet }}

{% endfor %}
#endif
""")

_fsm_src_template = Template("""\

#include <malloc.h>

/* ***************
 * Include Files *
 * ***************/

{% for inc_file in src.include -%}
#include "{{ inc_file }}"
{%- endfor %}

/* ***********
 * Functions *
 * ***********/

{% for code_snippet in src.code -%}
{{ code_snippet }}

{% endfor %}
""")

_fsm_decl_template = Template("""\
/* function options (EDIT) */
typedef struct {{ param.fopts.type }} {
    /* define your options struct here */
} {{ param.fopts.type }};

/* transition check */
typedef enum e{{ param.type }}Check {
\tE{{ prefix|upper }}_TR_RETREAT,
\tE{{ prefix|upper }}_TR_ADVANCE,
\tE{{ prefix|upper }}_TR_CONTINUE,
\tE{{ prefix|upper }}_TR_BADINPUT
} e{{ param.type }}Check;

/* states (enum) */
typedef enum e{{ param.type }}State {
{%- for state in param.states %}
\tE{{ prefix|upper }}_ST_{{ state|upper }},
{%- endfor %}
\tE{{ prefix|upper }}_NUM_STATES
} e{{ param.type }}State;

/* inputs (enum) */
typedef enum e{{ param.type }}Input {
{%- for input in param.inputs %}
\tE{{ prefix|upper }}_IN_{{ input|upper }},
{%- endfor %}
\tE{{ prefix|upper }}_NUM_INPUTS,
\tE{{ prefix|upper }}_NOINPUT
} e{{ param.type }}Input;

/* finite state machine struct */
typedef struct {{ param.type }} {
\te{{ param.type }}Input input;
\te{{ param.type }}Check check;
\te{{ param.type }}State cur;
\te{{ param.type }}State cmd;
\te{{ param.type }}State **transition_table;
\tvoid (***state_transitions)(struct {{ param.type }} *, {{ param.fopts.type }} *);
\tvoid (*run)(struct {{ param.type }} *, {{ param.fopts.type }} *, const e{{ param.type }}Input);
} {{ param.type }};

/* transition functions */
typedef void (*p{{ param.type }}StateTransitions)\
(struct {{ param.type }} *, {{ param.fopts.type }} *);

/* fsm declarations */
{%  for stateone in param.states -%}
{%- for statetwo in param.states -%}
{%- if stateone == statetwo or not stateone in param.transitiontable or statetwo in param.transitiontable[stateone] -%}
void {{ prefix|lower }}_{{ stateone|lower }}_{{ statetwo|lower }} \
({{ param.type }} *fsm, {{ param.fopts.type }} *{{ param.fopts.name }});
{%  endif -%}
{%  endfor -%}
{%  endfor -%}
void {{ prefix|lower }}_run ({{ param.type }} *fsm, \
{{ param.fopts.type }} *{{ param.fopts.name }}, \
const e{{ param.type }}Input input);

/* create */
{{ param.type }} *{{ prefix|lower }}_create(void);

/* free */
void {{ prefix|lower }}_free ({{ param.type }} *fsm);

""")

_fsm_fcns_template = Template("""\
/*--------------------------*
 *  RUNNING STATE FUNCTIONS *
 *--------------------------*/
{% for state in param.states %}
/**
 *  @brief Running function in state `{{ state|lower }}`
 */
void {{ prefix|lower }}_{{ state|lower }}_{{ state|lower }} \
({{ param.type }} *fsm, {{ param.fopts.type }} *{{ param.fopts.name }}) {

    /* check if this function was called from a transition. */
    if (fsm->check == E{{ prefix|upper }}_TR_ADVANCE) {
        /* transitioned from fsm->cur to fsm->cmd (here) */
    } else if (fsm->check == E{{ prefix|upper }}_TR_RETREAT) {
        /* fell back to fsm->cur (here) from fsm->cmd */
    } else {
        /* no prior transition (fsm->check == E{{ prefix|upper }}_TR_CONTINUE) */
    }

}
{% endfor %}

/*----------------------*
 * TRANSITION FUNCTIONS *
 *----------------------*/
{%  for stateone in param.states -%}{%  for statetwo in param.states -%}
{% if (stateone != statetwo) -%}
{% if not stateone in param.transitiontable or statetwo in param.transitiontable[stateone] %}
/**
 *  @brief Transition function from `{{ stateone|lower }}` to `{{ statetwo|lower }}`
 *
 *  @details
 *
 *  To advance to '{{ statetwo|lower }}' state, set
 *  `fsm->check = E{{ prefix|upper }}_TR_ADVANCE;`
 *  To return to '{{ stateone|lower }}' state, set
 *  `fsm->check = E{{ prefix|upper }}_TR_RETREAT;`
 */
void {{ prefix|lower }}_{{ stateone|lower }}_{{ statetwo|lower }} \
({{ param.type }} *fsm, {{ param.fopts.type }} *{{ param.fopts.name }}) {

    /* by default, do not transition (guard/retreat) */
    (void)({{ param.fopts.name }});
    fsm->check = E{{ prefix|upper }}_TR_RETREAT;

    /* Transition code goes here.

       NOTE: Before returning from this function,
       Consider setting transition to
       advance: fsm->check = E{{ prefix|upper }}_TR_ADVANCE; */

}
{% endif -%}
{% endif -%}
{% endfor -%}
{% endfor %}

/*-------------------*
 * RUN STATE MACHINE *
 *-------------------*/

/**
 *  @brief Run state machine
 */
void {{ prefix|lower }}_run ({{ param.type }} *fsm, \
{{ param.fopts.type }} *{{ param.fopts.name }}, \
const e{{ param.type }}Input input) {

    /* transition table - get command from input */
    if (input < E{{ prefix|upper }}_NUM_INPUTS) {
        fsm->input = input;
        fsm->cmd = fsm->transition_table[fsm->cur][input];
        if (fsm->cmd == fsm->cur) {
            /* not able to go to new input */
            fsm->check = E{{ prefix|upper }}_TR_BADINPUT;
        }
    }

    /* run process */
    if (fsm->state_transitions[fsm->cur][fsm->cmd] == NULL) {
        fsm->check = E{{ prefix|upper }}_TR_RETREAT;
    } else {
        fsm->state_transitions[fsm->cur][fsm->cmd](fsm, {{ param.fopts.name }});
    }

    /* advance to requested state or return to current state */
    if (fsm->cmd != fsm->cur) {
        if (fsm->check == E{{ prefix|upper }}_TR_ADVANCE) {
            fsm->state_transitions[fsm->cmd][fsm->cmd](fsm, {{ param.fopts.name }});
            fsm->cur = fsm->cmd;
        } else {
            fsm->state_transitions[fsm->cur][fsm->cur](fsm, {{ param.fopts.name }});
            fsm->cmd = fsm->cur;
        }
    }

    /* continue running */
    fsm->input = E{{ prefix|upper }}_NOINPUT;
    fsm->check = E{{ prefix|upper }}_TR_CONTINUE;
}

/*----------------------*
 * CREATE STATE MACHINE *
 *----------------------*/

/**
 *  @brief Create state machine
 */
{{ param.type }} * {{ prefix|lower }}_create (void) {

    {{ param.type }} *fsm = ({{ param.type }} *) malloc(sizeof({{ param.type }}));
    if (fsm == NULL) {
        return NULL;
    }

    fsm->input = E{{ prefix|upper }}_NOINPUT;
    fsm->check = E{{ prefix|upper }}_TR_CONTINUE;
    fsm->cur = E{{ prefix|upper }}_ST_{{ param.states|first|upper }};
    fsm->cmd = E{{ prefix|upper }}_ST_{{ param.states|first|upper }};
    fsm->run = {{ prefix|lower }}_run;

    // set future pointer allocations to NULL
    fsm->transition_table = NULL;
    fsm->state_transitions = NULL;

    /* transition table */

    fsm->transition_table = (e{{ param.type }}State **) malloc(E{{ prefix|upper }}_NUM_STATES * sizeof(e{{ param.type }}State *));
    if (fsm->transition_table == NULL) {
        {{ prefix|lower }}_free(fsm);
        return NULL;
    }
    // set future pointer allocations to NULL
    for (int k = 0; k < E{{ prefix|upper }}_NUM_STATES; k++) {
        fsm->transition_table[k] = NULL;
    }
    {%-  for state in param.states %}{% set stateloop = loop %}

    fsm->transition_table[{{ stateloop.index - 1 }}] = (e{{ param.type }}State *) malloc(E{{ prefix|upper }}_NUM_INPUT *, sizeof(e{{ param.type }}State));
    if (fsm->transition_table[{{ stateloop.index - 1 }}] == NULL) {
        {{ prefix|lower }}_free(fsm);
        return NULL;
    }
    {%  for next_state in param.transitiontable[state] %}
    {% if next_state in param.states -%}
    fsm->transition_table[{{ stateloop.index - 1 }}][{{ loop.index - 1 }}] = E{{ prefix|upper }}_ST_{{ next_state|upper }};
    {%- else -%}
    fsm->transition_table[{{ stateloop.index - 1 }}][{{ loop.index - 1 }}] = E{{ prefix|upper }}_ST_{{ state|upper }};
    {%- endif -%}
    {%- endfor %}
    {%- endfor %}

    /* state transitions */

    fsm->state_transitions = (p{{ param.type }}StateTransitions **) malloc(E{{ prefix|upper }}_NUM_STATES * sizeof(p{{ param.type }}StateTransitions *));
    if (fsm->state_transitions == NULL) {
        {{ prefix|lower }}_free(fsm);
        return NULL;
    }
    // set future pointer allocations to NULL
    for (int k = 0; k < E{{ prefix|upper }}_NUM_STATES; k++) {
        fsm->state_transitions[k] = NULL;
    }
    {%-  for stateone in param.states %}{% set stateoneloop = loop %}

    fsm->state_transitions[{{ stateoneloop.index - 1 }}] = (p{{ param.type }}StateTransitions *) malloc(E{{ prefix|upper }}_NUM_STATES * sizeof(p{{ param.type }}StateTransitions));
    if (fsm->state_transitions[{{ stateoneloop.index - 1 }}] == NULL) {
        {{ prefix|lower }}_free(fsm);
        return NULL;
    }
    {% for statetwo in param.states %}
    {% if not stateone == statetwo and stateone in param.transitiontable and not statetwo in param.transitiontable[stateone] -%}
    fsm->state_transitions[{{ stateoneloop.index - 1 }}][{{ loop.index - 1 }}] = NULL;
    {%- else -%}
    fsm->state_transitions[{{ stateoneloop.index - 1 }}][{{ loop.index - 1 }}] = {{ prefix|lower }}_{{ stateone|lower }}_{{ statetwo|lower }};
    {%- endif -%}
    {%- endfor %}
    {%- endfor %}

    return fsm;

}

/**
 *  @brief Free state machine
 */
void {{ prefix|lower }}_free ({{ param.type }} *fsm) {
    if (fsm != NULL) {
        if (fsm->transition_table != NULL) {
            for (int k = 0; k < E{{ prefix|upper }}_NUM_STATES; ++k)
            {
                if (fsm->transition_table[k] != NULL) {
                    free(fsm->transition_table[k]);
                    fsm->transition_table[k] = NULL;
                }
            }
            free(fsm->transition_table);
            fsm->transition_table = NULL;
        }
        if (fsm->state_transitions != NULL) {
            for (int k = 0; k < E{{ prefix|upper }}_NUM_STATES; ++k)
            {
                if (fsm->state_transitions[k] != NULL) {
                    free(fsm->state_transitions[k]);
                    fsm->state_transitions[k] = NULL;
                }
            }
            free(fsm->state_transitions);
            fsm->state_transitions = NULL;
        }
    }
    free(fsm);
}
""")

def mkdir_p(path):
    """Helper function to mimic bash command 'mkdir -p'"""
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

class Fsm(object):
    def __init__(self, fsm_param):
        """FSM parameters"""

        self.param = {
            'type': fsm_param['type'],
            'states': fsm_param['states'],
            'inputs': fsm_param['inputs'],
            'transitiontable': fsm_param['transitiontable'],
            'fopts': {
                'type': fsm_param['type'] + 'Opts',
                'name': 'fopts'
            }
        }

        if 'fopts' in fsm_param:
            if 'type' in fsm_param['fopts']:
                self.param['fopts']['type'] = fsm_param['fopts']['type']
            if 'name' in fsm_param['fopts']:
                self.param['fopts']['name'] = fsm_param['fopts']['name']

    def genccode(self, folder, prefix):
        """Create C code"""

        # render code templates
        fsm_decl = _fsm_decl_template.render(prefix=prefix,
                                             param=self.param)
        fsm_fcns = _fsm_fcns_template.render(prefix=prefix,
                                             param=self.param)

        # include file
        header = {
            'filename': "{b}.h".format(b=prefix),
            'basename': prefix,
            'code': [fsm_decl],
            'include': []
        }

        # src file
        src = {
            'filename': "{b}.c".format(b=prefix),
            'code': [fsm_fcns],
            'include': [header['filename']]
        }

        # code wrapper templates
        header_str = _fsm_header_template.render(header=header)
        src_str = _fsm_src_template.render(src=src)

        # make directory
        path = os.path.abspath(folder)
        mkdir_p(path)

        # write header file
        f = open(os.path.join(path, header['filename']), "w")
        f.write(header_str)
        f.close()

        # write source file
        f = open(os.path.join(path, src['filename']), "w")
        f.write(src_str)
        f.close()
<commit_msg>Fix transition table malloc for inputs<commit_after>import os
import errno
from jinja2 import Template

_fsm_header_template = Template("""\
#ifndef {{ header.basename|upper }}_H
#define {{ header.basename|upper }}_H

/* ***************************
 * Typedefs and Declarations *
 * ***************************/

{% for code_snippet in header.code -%}
{{ code_snippet }}

{% endfor %}
#endif
""")

_fsm_src_template = Template("""\

#include <malloc.h>

/* ***************
 * Include Files *
 * ***************/

{% for inc_file in src.include -%}
#include "{{ inc_file }}"
{%- endfor %}

/* ***********
 * Functions *
 * ***********/

{% for code_snippet in src.code -%}
{{ code_snippet }}

{% endfor %}
""")

_fsm_decl_template = Template("""\
/* function options (EDIT) */
typedef struct {{ param.fopts.type }} {
    /* define your options struct here */
} {{ param.fopts.type }};

/* transition check */
typedef enum e{{ param.type }}Check {
\tE{{ prefix|upper }}_TR_RETREAT,
\tE{{ prefix|upper }}_TR_ADVANCE,
\tE{{ prefix|upper }}_TR_CONTINUE,
\tE{{ prefix|upper }}_TR_BADINPUT
} e{{ param.type }}Check;

/* states (enum) */
typedef enum e{{ param.type }}State {
{%- for state in param.states %}
\tE{{ prefix|upper }}_ST_{{ state|upper }},
{%- endfor %}
\tE{{ prefix|upper }}_NUM_STATES
} e{{ param.type }}State;

/* inputs (enum) */
typedef enum e{{ param.type }}Input {
{%- for input in param.inputs %}
\tE{{ prefix|upper }}_IN_{{ input|upper }},
{%- endfor %}
\tE{{ prefix|upper }}_NUM_INPUTS,
\tE{{ prefix|upper }}_NOINPUT
} e{{ param.type }}Input;

/* finite state machine struct */
typedef struct {{ param.type }} {
\te{{ param.type }}Input input;
\te{{ param.type }}Check check;
\te{{ param.type }}State cur;
\te{{ param.type }}State cmd;
\te{{ param.type }}State **transition_table;
\tvoid (***state_transitions)(struct {{ param.type }} *, {{ param.fopts.type }} *);
\tvoid (*run)(struct {{ param.type }} *, {{ param.fopts.type }} *, const e{{ param.type }}Input);
} {{ param.type }};

/* transition functions */
typedef void (*p{{ param.type }}StateTransitions)\
(struct {{ param.type }} *, {{ param.fopts.type }} *);

/* fsm declarations */
{%  for stateone in param.states -%}
{%- for statetwo in param.states -%}
{%- if stateone == statetwo or not stateone in param.transitiontable or statetwo in param.transitiontable[stateone] -%}
void {{ prefix|lower }}_{{ stateone|lower }}_{{ statetwo|lower }} \
({{ param.type }} *fsm, {{ param.fopts.type }} *{{ param.fopts.name }});
{%  endif -%}
{%  endfor -%}
{%  endfor -%}
void {{ prefix|lower }}_run ({{ param.type }} *fsm, \
{{ param.fopts.type }} *{{ param.fopts.name }}, \
const e{{ param.type }}Input input);

/* create */
{{ param.type }} *{{ prefix|lower }}_create(void);

/* free */
void {{ prefix|lower }}_free ({{ param.type }} *fsm);

""")

_fsm_fcns_template = Template("""\
/*--------------------------*
 *  RUNNING STATE FUNCTIONS *
 *--------------------------*/
{% for state in param.states %}
/**
 *  @brief Running function in state `{{ state|lower }}`
 */
void {{ prefix|lower }}_{{ state|lower }}_{{ state|lower }} \
({{ param.type }} *fsm, {{ param.fopts.type }} *{{ param.fopts.name }}) {

    /* check if this function was called from a transition. */
    if (fsm->check == E{{ prefix|upper }}_TR_ADVANCE) {
        /* transitioned from fsm->cur to fsm->cmd (here) */
    } else if (fsm->check == E{{ prefix|upper }}_TR_RETREAT) {
        /* fell back to fsm->cur (here) from fsm->cmd */
    } else {
        /* no prior transition (fsm->check == E{{ prefix|upper }}_TR_CONTINUE) */
    }

}
{% endfor %}

/*----------------------*
 * TRANSITION FUNCTIONS *
 *----------------------*/
{%  for stateone in param.states -%}{%  for statetwo in param.states -%}
{% if (stateone != statetwo) -%}
{% if not stateone in param.transitiontable or statetwo in param.transitiontable[stateone] %}
/**
 *  @brief Transition function from `{{ stateone|lower }}` to `{{ statetwo|lower }}`
 *
 *  @details
 *
 *  To advance to '{{ statetwo|lower }}' state, set
 *  `fsm->check = E{{ prefix|upper }}_TR_ADVANCE;`
 *  To return to '{{ stateone|lower }}' state, set
 *  `fsm->check = E{{ prefix|upper }}_TR_RETREAT;`
 */
void {{ prefix|lower }}_{{ stateone|lower }}_{{ statetwo|lower }} \
({{ param.type }} *fsm, {{ param.fopts.type }} *{{ param.fopts.name }}) {

    /* by default, do not transition (guard/retreat) */
    (void)({{ param.fopts.name }});
    fsm->check = E{{ prefix|upper }}_TR_RETREAT;

    /* Transition code goes here.

       NOTE: Before returning from this function,
       Consider setting transition to
       advance: fsm->check = E{{ prefix|upper }}_TR_ADVANCE; */

}
{% endif -%}
{% endif -%}
{% endfor -%}
{% endfor %}

/*-------------------*
 * RUN STATE MACHINE *
 *-------------------*/

/**
 *  @brief Run state machine
 */
void {{ prefix|lower }}_run ({{ param.type }} *fsm, \
{{ param.fopts.type }} *{{ param.fopts.name }}, \
const e{{ param.type }}Input input) {

    /* transition table - get command from input */
    if (input < E{{ prefix|upper }}_NUM_INPUTS) {
        fsm->input = input;
        fsm->cmd = fsm->transition_table[fsm->cur][input];
        if (fsm->cmd == fsm->cur) {
            /* not able to go to new input */
            fsm->check = E{{ prefix|upper }}_TR_BADINPUT;
        }
    }

    /* run process */
    if (fsm->state_transitions[fsm->cur][fsm->cmd] == NULL) {
        fsm->check = E{{ prefix|upper }}_TR_RETREAT;
    } else {
        fsm->state_transitions[fsm->cur][fsm->cmd](fsm, {{ param.fopts.name }});
    }

    /* advance to requested state or return to current state */
    if (fsm->cmd != fsm->cur) {
        if (fsm->check == E{{ prefix|upper }}_TR_ADVANCE) {
            fsm->state_transitions[fsm->cmd][fsm->cmd](fsm, {{ param.fopts.name }});
            fsm->cur = fsm->cmd;
        } else {
            fsm->state_transitions[fsm->cur][fsm->cur](fsm, {{ param.fopts.name }});
            fsm->cmd = fsm->cur;
        }
    }

    /* continue running */
    fsm->input = E{{ prefix|upper }}_NOINPUT;
    fsm->check = E{{ prefix|upper }}_TR_CONTINUE;
}

/*----------------------*
 * CREATE STATE MACHINE *
 *----------------------*/

/**
 *  @brief Create state machine
 */
{{ param.type }} * {{ prefix|lower }}_create (void) {

    {{ param.type }} *fsm = ({{ param.type }} *) malloc(sizeof({{ param.type }}));
    if (fsm == NULL) {
        return NULL;
    }

    fsm->input = E{{ prefix|upper }}_NOINPUT;
    fsm->check = E{{ prefix|upper }}_TR_CONTINUE;
    fsm->cur = E{{ prefix|upper }}_ST_{{ param.states|first|upper }};
    fsm->cmd = E{{ prefix|upper }}_ST_{{ param.states|first|upper }};
    fsm->run = {{ prefix|lower }}_run;

    // set future pointer allocations to NULL
    fsm->transition_table = NULL;
    fsm->state_transitions = NULL;

    /* transition table */

    fsm->transition_table = (e{{ param.type }}State **) malloc(E{{ prefix|upper }}_NUM_STATES * sizeof(e{{ param.type }}State *));
    if (fsm->transition_table == NULL) {
        {{ prefix|lower }}_free(fsm);
        return NULL;
    }
    // set future pointer allocations to NULL
    for (int k = 0; k < E{{ prefix|upper }}_NUM_STATES; k++) {
        fsm->transition_table[k] = NULL;
    }
    {%-  for state in param.states %}{% set stateloop = loop %}

    fsm->transition_table[{{ stateloop.index - 1 }}] = (e{{ param.type }}State *) malloc(E{{ prefix|upper }}_NUM_INPUTS * sizeof(e{{ param.type }}State));
    if (fsm->transition_table[{{ stateloop.index - 1 }}] == NULL) {
        {{ prefix|lower }}_free(fsm);
        return NULL;
    }
    {%  for next_state in param.transitiontable[state] %}
    {% if next_state in param.states -%}
    fsm->transition_table[{{ stateloop.index - 1 }}][{{ loop.index - 1 }}] = E{{ prefix|upper }}_ST_{{ next_state|upper }};
    {%- else -%}
    fsm->transition_table[{{ stateloop.index - 1 }}][{{ loop.index - 1 }}] = E{{ prefix|upper }}_ST_{{ state|upper }};
    {%- endif -%}
    {%- endfor %}
    {%- endfor %}

    /* state transitions */

    fsm->state_transitions = (p{{ param.type }}StateTransitions **) malloc(E{{ prefix|upper }}_NUM_STATES * sizeof(p{{ param.type }}StateTransitions *));
    if (fsm->state_transitions == NULL) {
        {{ prefix|lower }}_free(fsm);
        return NULL;
    }
    // set future pointer allocations to NULL
    for (int k = 0; k < E{{ prefix|upper }}_NUM_STATES; k++) {
        fsm->state_transitions[k] = NULL;
    }
    {%-  for stateone in param.states %}{% set stateoneloop = loop %}

    fsm->state_transitions[{{ stateoneloop.index - 1 }}] = (p{{ param.type }}StateTransitions *) malloc(E{{ prefix|upper }}_NUM_STATES * sizeof(p{{ param.type }}StateTransitions));
    if (fsm->state_transitions[{{ stateoneloop.index - 1 }}] == NULL) {
        {{ prefix|lower }}_free(fsm);
        return NULL;
    }
    {% for statetwo in param.states %}
    {% if not stateone == statetwo and stateone in param.transitiontable and not statetwo in param.transitiontable[stateone] -%}
    fsm->state_transitions[{{ stateoneloop.index - 1 }}][{{ loop.index - 1 }}] = NULL;
    {%- else -%}
    fsm->state_transitions[{{ stateoneloop.index - 1 }}][{{ loop.index - 1 }}] = {{ prefix|lower }}_{{ stateone|lower }}_{{ statetwo|lower }};
    {%- endif -%}
    {%- endfor %}
    {%- endfor %}

    return fsm;

}

/**
 *  @brief Free state machine
 */
void {{ prefix|lower }}_free ({{ param.type }} *fsm) {
    if (fsm != NULL) {
        if (fsm->transition_table != NULL) {
            for (int k = 0; k < E{{ prefix|upper }}_NUM_STATES; ++k)
            {
                if (fsm->transition_table[k] != NULL) {
                    free(fsm->transition_table[k]);
                    fsm->transition_table[k] = NULL;
                }
            }
            free(fsm->transition_table);
            fsm->transition_table = NULL;
        }
        if (fsm->state_transitions != NULL) {
            for (int k = 0; k < E{{ prefix|upper }}_NUM_STATES; ++k)
            {
                if (fsm->state_transitions[k] != NULL) {
                    free(fsm->state_transitions[k]);
                    fsm->state_transitions[k] = NULL;
                }
            }
            free(fsm->state_transitions);
            fsm->state_transitions = NULL;
        }
    }
    free(fsm);
}
""")

def mkdir_p(path):
    """Helper function to mimic bash command 'mkdir -p'"""
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

class Fsm(object):
    def __init__(self, fsm_param):
        """FSM parameters"""

        self.param = {
            'type': fsm_param['type'],
            'states': fsm_param['states'],
            'inputs': fsm_param['inputs'],
            'transitiontable': fsm_param['transitiontable'],
            'fopts': {
                'type': fsm_param['type'] + 'Opts',
                'name': 'fopts'
            }
        }

        if 'fopts' in fsm_param:
            if 'type' in fsm_param['fopts']:
                self.param['fopts']['type'] = fsm_param['fopts']['type']
            if 'name' in fsm_param['fopts']:
                self.param['fopts']['name'] = fsm_param['fopts']['name']

    def genccode(self, folder, prefix):
        """Create C code"""

        # render code templates
        fsm_decl = _fsm_decl_template.render(prefix=prefix,
                                             param=self.param)
        fsm_fcns = _fsm_fcns_template.render(prefix=prefix,
                                             param=self.param)

        # include file
        header = {
            'filename': "{b}.h".format(b=prefix),
            'basename': prefix,
            'code': [fsm_decl],
            'include': []
        }

        # src file
        src = {
            'filename': "{b}.c".format(b=prefix),
            'code': [fsm_fcns],
            'include': [header['filename']]
        }

        # code wrapper templates
        header_str = _fsm_header_template.render(header=header)
        src_str = _fsm_src_template.render(src=src)

        # make directory
        path = os.path.abspath(folder)
        mkdir_p(path)

        # write header file
        f = open(os.path.join(path, header['filename']), "w")
        f.write(header_str)
        f.close()

        # write source file
        f = open(os.path.join(path, src['filename']), "w")
        f.write(src_str)
        f.close()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
@author: Federico Cerchiari <federicocerchiari@gmail.com>
Places used by Tempy to choose the right TempyREPR to use.
Magically created starting from the tags module.
"""
import importlib
from .tempyrepr import TempyPlace

class Inside(TempyPlace):
    def _reprscore_container_parent(self, container, child):
        if container.parent.__class__ == self._pointer_class:
            return 1
        return 0

class SiblingLooker(TempyPlace):
    @classmethod
    def _check_container_siblings(cls, container, start=-1, stop=2):
        container_index = cls._content_index(container.parent, container)
        before = container.parent.childs[container_index+start:container_index]
        after = container.parent.childs[container_index+1:container_index+stop]
        for sibling in  before + after:
            if sibling.__class__ == cls._pointer_class:
                return 1
        return 0

class Near(SiblingLooker):
    def _reprscore_container_siblings(self, container, child):
        return self._check_container_siblings(container)

class Before(SiblingLooker):
    def _reprscore_container_following(self, container, child):
        return self._check_container_siblings(container, start=0)

class After(SiblingLooker):
    def _reprscore_container_previous(self, container, child):
        return self._check_container_siblings(container, stop=0)

# Creation of a TempyPlace class for every tag defined in the tags.py module
all_tags = importlib.import_module('.tags', package='tempy')
for tag in dir(all_tags):
    if not tag.startswith('__'):
        for place_type in (Inside, Near, Before, After):
            place_tag_name = '%s%s' % (place_type.__name__, tag)
            place_cls = type(place_tag_name, (place_type, ), {'_pointer_class': getattr(all_tags, tag), '_base_place': False})
            # We put the new dynamically created class inside locals to make it avaiable from the outside
            locals()[place_tag_name] = place_cls
<commit_msg>pep8 + docstrings. [skip ci]<commit_after># -*- coding: utf-8 -*-
"""
@author: Federico Cerchiari <federicocerchiari@gmail.com>
Places used by Tempy to choose the right TempyREPR to use.
Magically created starting from the tags module.
"""
import importlib
from .tempyrepr import TempyPlace

class Inside(TempyPlace):
    """Check if a TempyREPR object's container is inside a certain Tempy Tag"""
    def _reprscore_container_parent(self, container, child):
        if container.parent.__class__ == self._pointer_class:
            return 1
        return 0

class Sibling(TempyPlace):
    """Base class for TempyPlaces that depends on the TempyREPR object's container siblings"""
    @classmethod
    def _check_container_siblings(cls, container, start=-1, stop=2):
        container_index = cls._content_index(container.parent, container)
        before = container.parent.childs[container_index+start:container_index]
        after = container.parent.childs[container_index+1:container_index+stop]
        for sibling in before + after:
            if sibling.__class__ == cls._pointer_class:
                return 1
        return 0

class Near(Sibling):
    """Check if a TempyREPR object's container if near a certain Tempy Tag"""
    def _reprscore_container_siblings(self, container, child):
        return self._check_container_siblings(container)

class Before(Sibling):
    """Check if a TempyREPR object's container if before a certain Tempy Tag"""
    def _reprscore_container_following(self, container, child):
        return self._check_container_siblings(container, start=0)

class After(Sibling):
    """Check if a TempyREPR object's container if after a certain Tempy Tag"""
    def _reprscore_container_previous(self, container, child):
        return self._check_container_siblings(container, stop=0)

# Creation of a TempyPlace class for every tag defined in the tags.py module
all_tags = importlib.import_module('.tags', package='tempy')
for tag in dir(all_tags):
    if not tag.startswith('__'):
        for place_type in (Inside, Near, Before, After):
            place_tag_name = '%s%s' % (place_type.__name__, tag)
            # Dynamic class definition
            place_cls = type(place_tag_name, (place_type, ), {'_pointer_class': getattr(all_tags, tag), '_base_place': False})
            # We put the new dynamically created class inside locals to make it avaiable from the outside
            locals()[place_tag_name] = place_cls
<|endoftext|>


<commit_before>"""
Django Forms for interacting with Part objects
"""

# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django import forms
from django.utils.translation import ugettext_lazy as _

from mptt.fields import TreeNodeChoiceField

from InvenTree.forms import HelperForm
from InvenTree.helpers import GetExportFormats, clean_decimal
from InvenTree.fields import RoundingDecimalFormField

import common.models
from common.forms import MatchItemForm

from .models import Part, PartCategory, PartRelated
from .models import PartParameterTemplate, PartParameter
from .models import PartCategoryParameterTemplate
from .models import PartSellPriceBreak, PartInternalPriceBreak

class PartModelChoiceField(forms.ModelChoiceField):
    """ Extending string representation of Part instance with available stock """

    def label_from_instance(self, part):

        label = str(part)

        # Optionally display available part quantity
        if common.models.InvenTreeSetting.get_setting('PART_SHOW_QUANTITY_IN_FORMS'):
            label += f" - {part.available_stock}"

        return label

class PartImageDownloadForm(HelperForm):
    """
    Form for downloading an image from a URL
    """

    url = forms.URLField(
        label=_('URL'),
        help_text=_('Image URL'),
        required=True,
    )

    class Meta:
        model = Part
        fields = [
            'url',
        ]

class BomExportForm(forms.Form):
    """ Simple form to let user set BOM export options,
    before exporting a BOM (bill of materials) file.
    """

    file_format = forms.ChoiceField(label=_("File Format"), help_text=_("Select output file format"))

    cascading = forms.BooleanField(label=_("Cascading"), required=False, initial=True, help_text=_("Download cascading / multi-level BOM"))

    levels = forms.IntegerField(label=_("Levels"), required=True, initial=0, help_text=_("Select maximum number of BOM levels to export (0 = all levels)"))

    parameter_data = forms.BooleanField(label=_("Include Parameter Data"), required=False, initial=False, help_text=_("Include part parameters data in exported BOM"))

    stock_data = forms.BooleanField(label=_("Include Stock Data"), required=False, initial=False, help_text=_("Include part stock data in exported BOM"))

    manufacturer_data = forms.BooleanField(label=_("Include Manufacturer Data"), required=False, initial=True, help_text=_("Include part manufacturer data in exported BOM"))

    supplier_data = forms.BooleanField(label=_("Include Supplier Data"), required=False, initial=True, help_text=_("Include part supplier data in exported BOM"))

    def get_choices(self):
        """ BOM export format choices """

        return [(x, x.upper()) for x in GetExportFormats()]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.fields['file_format'].choices = self.get_choices()

class BomDuplicateForm(HelperForm):
    """
    Simple confirmation form for BOM duplication.

    Select which parent to select from.
    """

    parent = PartModelChoiceField(
        label=_('Parent Part'),
        help_text=_('Select parent part to copy BOM from'),
        queryset=Part.objects.filter(is_template=True),
    )

    clear = forms.BooleanField(
        required=False, initial=True,
        help_text=_('Clear existing BOM items')
    )

    confirm = forms.BooleanField(
        required=False, initial=False,
        label=_('Confirm'),
        help_text=_('Confirm BOM duplication')
    )

    class Meta:
        model = Part
        fields = [
            'parent',
            'clear',
            'confirm',
        ]

class BomValidateForm(HelperForm):
    """ Simple confirmation form for BOM validation.
    User is presented with a single checkbox input,
    to confirm that the BOM for this part is valid
    """

    validate = forms.BooleanField(required=False, initial=False, label=_('validate'), help_text=_('Confirm that the BOM is correct'))

    class Meta:
        model = Part
        fields = [
            'validate'
        ]

class BomMatchItemForm(MatchItemForm):
    """ Override MatchItemForm fields """

    def get_special_field(self, col_guess, row, file_manager):
        """ Set special fields """

        # set quantity field
        if 'quantity' in col_guess.lower():
            return forms.CharField(
                required=False,
                widget=forms.NumberInput(attrs={
                    'name': 'quantity' + str(row['index']),
                    'class': 'numberinput',
                    'type': 'number',
                    'min': '0',
                    'step': 'any',
                    'value': clean_decimal(row.get('quantity', '')),
                })
            )

        # return default
        return super().get_special_field(col_guess, row, file_manager)

class CreatePartRelatedForm(HelperForm):
    """ Form for creating a PartRelated object """

    class Meta:
        model = PartRelated
        fields = [
            'part_1',
            'part_2',
        ]
        labels = {
            'part_2': _('Related Part'),
        }

class SetPartCategoryForm(forms.Form):
    """ Form for setting the category of multiple Part objects """

    part_category = TreeNodeChoiceField(queryset=PartCategory.objects.all(), required=True, help_text=_('Select part category'))

class EditPartParameterTemplateForm(HelperForm):
    """ Form for editing a PartParameterTemplate object """

    class Meta:
        model = PartParameterTemplate
        fields = [
            'name',
            'units'
        ]

class EditPartParameterForm(HelperForm):
    """ Form for editing a PartParameter object """

    class Meta:
        model = PartParameter
        fields = [
            'part',
            'template',
            'data'
        ]

class EditCategoryForm(HelperForm):
    """ Form for editing a PartCategory object """

    field_prefix = {
        'default_keywords': 'fa-key',
    }

    class Meta:
        model = PartCategory
        fields = [
            'parent',
            'name',
            'description',
            'default_location',
            'default_keywords',
        ]

class EditCategoryParameterTemplateForm(HelperForm):
    """ Form for editing a PartCategoryParameterTemplate object """

    add_to_same_level_categories = forms.BooleanField(required=False,
                                                      initial=False,
                                                      help_text=_('Add parameter template to same level categories'))

    add_to_all_categories = forms.BooleanField(required=False,
                                               initial=False,
                                               help_text=_('Add parameter template to all categories'))

    class Meta:
        model = PartCategoryParameterTemplate
        fields = [
            'category',
            'parameter_template',
            'default_value',
            'add_to_same_level_categories',
            'add_to_all_categories',
        ]

class PartPriceForm(forms.Form):
    """ Simple form for viewing part pricing information """

    quantity = forms.IntegerField(
        required=True,
        initial=1,
        label=_('Quantity'),
        help_text=_('Input quantity for price calculation')
    )

    class Meta:
        model = Part
        fields = [
            'quantity',
        ]

class EditPartSalePriceBreakForm(HelperForm):
    """
    Form for creating / editing a sale price for a part
    """

    quantity = RoundingDecimalFormField(max_digits=10, decimal_places=5, label=_('Quantity'))

    class Meta:
        model = PartSellPriceBreak
        fields = [
            'part',
            'quantity',
            'price',
        ]

class EditPartInternalPriceBreakForm(HelperForm):
    """
    Form for creating / editing a internal price for a part
    """

    quantity = RoundingDecimalFormField(max_digits=10, decimal_places=5, label=_('Quantity'))

    class Meta:
        model = PartInternalPriceBreak
        fields = [
            'part',
            'quantity',
            'price',
        ]
<commit_msg>Cleanup unused form<commit_after>"""
Django Forms for interacting with Part objects
"""

# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django import forms
from django.utils.translation import ugettext_lazy as _

from mptt.fields import TreeNodeChoiceField

from InvenTree.forms import HelperForm
from InvenTree.helpers import GetExportFormats, clean_decimal
from InvenTree.fields import RoundingDecimalFormField

import common.models
from common.forms import MatchItemForm

from .models import Part, PartCategory, PartRelated
from .models import PartParameterTemplate, PartParameter
from .models import PartCategoryParameterTemplate
from .models import PartSellPriceBreak, PartInternalPriceBreak

class PartModelChoiceField(forms.ModelChoiceField):
    """ Extending string representation of Part instance with available stock """

    def label_from_instance(self, part):

        label = str(part)

        # Optionally display available part quantity
        if common.models.InvenTreeSetting.get_setting('PART_SHOW_QUANTITY_IN_FORMS'):
            label += f" - {part.available_stock}"

        return label

class PartImageDownloadForm(HelperForm):
    """
    Form for downloading an image from a URL
    """

    url = forms.URLField(
        label=_('URL'),
        help_text=_('Image URL'),
        required=True,
    )

    class Meta:
        model = Part
        fields = [
            'url',
        ]

class BomExportForm(forms.Form):
    """ Simple form to let user set BOM export options,
    before exporting a BOM (bill of materials) file.
    """

    file_format = forms.ChoiceField(label=_("File Format"), help_text=_("Select output file format"))

    cascading = forms.BooleanField(label=_("Cascading"), required=False, initial=True, help_text=_("Download cascading / multi-level BOM"))

    levels = forms.IntegerField(label=_("Levels"), required=True, initial=0, help_text=_("Select maximum number of BOM levels to export (0 = all levels)"))

    parameter_data = forms.BooleanField(label=_("Include Parameter Data"), required=False, initial=False, help_text=_("Include part parameters data in exported BOM"))

    stock_data = forms.BooleanField(label=_("Include Stock Data"), required=False, initial=False, help_text=_("Include part stock data in exported BOM"))

    manufacturer_data = forms.BooleanField(label=_("Include Manufacturer Data"), required=False, initial=True, help_text=_("Include part manufacturer data in exported BOM"))

    supplier_data = forms.BooleanField(label=_("Include Supplier Data"), required=False, initial=True, help_text=_("Include part supplier data in exported BOM"))

    def get_choices(self):
        """ BOM export format choices """

        return [(x, x.upper()) for x in GetExportFormats()]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.fields['file_format'].choices = self.get_choices()

class BomDuplicateForm(HelperForm):
    """
    Simple confirmation form for BOM duplication.

    Select which parent to select from.
    """

    parent = PartModelChoiceField(
        label=_('Parent Part'),
        help_text=_('Select parent part to copy BOM from'),
        queryset=Part.objects.filter(is_template=True),
    )

    clear = forms.BooleanField(
        required=False, initial=True,
        help_text=_('Clear existing BOM items')
    )

    confirm = forms.BooleanField(
        required=False, initial=False,
        label=_('Confirm'),
        help_text=_('Confirm BOM duplication')
    )

    class Meta:
        model = Part
        fields = [
            'parent',
            'clear',
            'confirm',
        ]

class BomValidateForm(HelperForm):
    """ Simple confirmation form for BOM validation.
    User is presented with a single checkbox input,
    to confirm that the BOM for this part is valid
    """

    validate = forms.BooleanField(required=False, initial=False, label=_('validate'), help_text=_('Confirm that the BOM is correct'))

    class Meta:
        model = Part
        fields = [
            'validate'
        ]

class BomMatchItemForm(MatchItemForm):
    """ Override MatchItemForm fields """

    def get_special_field(self, col_guess, row, file_manager):
        """ Set special fields """

        # set quantity field
        if 'quantity' in col_guess.lower():
            return forms.CharField(
                required=False,
                widget=forms.NumberInput(attrs={
                    'name': 'quantity' + str(row['index']),
                    'class': 'numberinput',
                    'type': 'number',
                    'min': '0',
                    'step': 'any',
                    'value': clean_decimal(row.get('quantity', '')),
                })
            )

        # return default
        return super().get_special_field(col_guess, row, file_manager)

class CreatePartRelatedForm(HelperForm):
    """ Form for creating a PartRelated object """

    class Meta:
        model = PartRelated
        fields = [
            'part_1',
            'part_2',
        ]
        labels = {
            'part_2': _('Related Part'),
        }

class SetPartCategoryForm(forms.Form):
    """ Form for setting the category of multiple Part objects """

    part_category = TreeNodeChoiceField(queryset=PartCategory.objects.all(), required=True, help_text=_('Select part category'))

class EditPartParameterTemplateForm(HelperForm):
    """ Form for editing a PartParameterTemplate object """

    class Meta:
        model = PartParameterTemplate
        fields = [
            'name',
            'units'
        ]

class EditCategoryForm(HelperForm):
    """ Form for editing a PartCategory object """

    field_prefix = {
        'default_keywords': 'fa-key',
    }

    class Meta:
        model = PartCategory
        fields = [
            'parent',
            'name',
            'description',
            'default_location',
            'default_keywords',
        ]

class EditCategoryParameterTemplateForm(HelperForm):
    """ Form for editing a PartCategoryParameterTemplate object """

    add_to_same_level_categories = forms.BooleanField(required=False,
                                                      initial=False,
                                                      help_text=_('Add parameter template to same level categories'))

    add_to_all_categories = forms.BooleanField(required=False,
                                               initial=False,
                                               help_text=_('Add parameter template to all categories'))

    class Meta:
        model = PartCategoryParameterTemplate
        fields = [
            'category',
            'parameter_template',
            'default_value',
            'add_to_same_level_categories',
            'add_to_all_categories',
        ]

class PartPriceForm(forms.Form):
    """ Simple form for viewing part pricing information """

    quantity = forms.IntegerField(
        required=True,
        initial=1,
        label=_('Quantity'),
        help_text=_('Input quantity for price calculation')
    )

    class Meta:
        model = Part
        fields = [
            'quantity',
        ]

class EditPartSalePriceBreakForm(HelperForm):
    """
    Form for creating / editing a sale price for a part
    """

    quantity = RoundingDecimalFormField(max_digits=10, decimal_places=5, label=_('Quantity'))

    class Meta:
        model = PartSellPriceBreak
        fields = [
            'part',
            'quantity',
            'price',
        ]

class EditPartInternalPriceBreakForm(HelperForm):
    """
    Form for creating / editing a internal price for a part
    """

    quantity = RoundingDecimalFormField(max_digits=10, decimal_places=5, label=_('Quantity'))

    class Meta:
        model = PartInternalPriceBreak
        fields = [
            'part',
            'quantity',
            'price',
        ]
<|endoftext|>