astropy.io.ascii.fastbasic のソースコード

# Licensed under a 3-clause BSD style license - see LICENSE.rst

import re
import copy
from collections import OrderedDict

from . import core
from astropy.table import Table
from . import cparser
from astropy.utils.misc import _set_locale


class FastBasic(metaclass=core.MetaBaseReader):
    """
    This class is intended to handle the same format addressed by the
    ordinary :class:`Basic` writer, but it acts as a wrapper for underlying C
    code and is therefore much faster. Unlike the other ASCII readers and
    writers, this class is not very extensible and is restricted
    by optimization requirements.
    """
    _format_name = 'fast_basic'
    _description = 'Basic table with custom delimiter using the fast C engine'
    _fast = True
    fill_extra_cols = False
    guessing = False
    strict_names = False

    def __init__(self, default_kwargs={}, **user_kwargs):
        # Make sure user does not set header_start to None for a reader
        # that expects a non-None value (i.e. a number >= 0).  This mimics
        # what happens in the Basic reader.
        if (default_kwargs.get('header_start', 0) is not None
                and user_kwargs.get('header_start', 0) is None):
            raise ValueError('header_start cannot be set to None for this Reader')

        # Set up kwargs and copy any user kwargs.  Use deepcopy user kwargs
        # since they may contain a dict item which would end up as a ref to the
        # original and get munged later (e.g. in cparser.pyx validation of
        # fast_reader dict).
        kwargs = copy.deepcopy(default_kwargs)
        kwargs.update(copy.deepcopy(user_kwargs))

        delimiter = kwargs.pop('delimiter', ' ')
        self.delimiter = str(delimiter) if delimiter is not None else None
        self.write_comment = kwargs.get('comment', '# ')
        self.comment = kwargs.pop('comment', '#')
        if self.comment is not None:
            self.comment = str(self.comment)
        self.quotechar = str(kwargs.pop('quotechar', '"'))
        self.header_start = kwargs.pop('header_start', 0)
        # If data_start is not specified, start reading
        # data right after the header line
        data_start_default = user_kwargs.get('data_start', self.header_start
                                             + 1 if self.header_start is not None else 1)
        self.data_start = kwargs.pop('data_start', data_start_default)
        self.kwargs = kwargs
        self.strip_whitespace_lines = True
        self.strip_whitespace_fields = True

    def _read_header(self):
        # Use the tokenizer by default -- this method
        # can be overridden for specialized headers
        self.engine.read_header()

    def read(self, table):
        """
        Read input data (file-like object, filename, list of strings, or
        single string) into a Table and return the result.
        """
        if self.comment is not None and len(self.comment) != 1:
            raise core.ParameterError("The C reader does not support a comment regex")
        elif self.data_start is None:
            raise core.ParameterError("The C reader does not allow data_start to be None")
        elif self.header_start is not None and self.header_start < 0 and \
                not isinstance(self, FastCommentedHeader):
            raise core.ParameterError("The C reader does not allow header_start to be "
                                      "negative except for commented-header files")
        elif self.data_start < 0:
            raise core.ParameterError("The C reader does not allow data_start to be negative")
        elif len(self.delimiter) != 1:
            raise core.ParameterError("The C reader only supports 1-char delimiters")
        elif len(self.quotechar) != 1:
            raise core.ParameterError("The C reader only supports a length-1 quote character")
        elif 'converters' in self.kwargs:
            raise core.ParameterError("The C reader does not support passing "
                                      "specialized converters")
        elif 'encoding' in self.kwargs:
            raise core.ParameterError("The C reader does not use the encoding parameter")
        elif 'Outputter' in self.kwargs:
            raise core.ParameterError("The C reader does not use the Outputter parameter")
        elif 'Inputter' in self.kwargs:
            raise core.ParameterError("The C reader does not use the Inputter parameter")
        elif 'data_Splitter' in self.kwargs or 'header_Splitter' in self.kwargs:
            raise core.ParameterError("The C reader does not use a Splitter class")

        self.strict_names = self.kwargs.pop('strict_names', False)

        # Process fast_reader kwarg, which may or may not exist (though ui.py will always
        # pass this as a dict with at least 'enable' set).
        fast_reader = self.kwargs.get('fast_reader', True)
        if not isinstance(fast_reader, dict):
            fast_reader = {}

        fast_reader.pop('enable', None)
        self.return_header_chars = fast_reader.pop('return_header_chars', False)
        # Put fast_reader dict back into kwargs.
        self.kwargs['fast_reader'] = fast_reader

        self.engine = cparser.CParser(table, self.strip_whitespace_lines,
                                      self.strip_whitespace_fields,
                                      delimiter=self.delimiter,
                                      header_start=self.header_start,
                                      comment=self.comment,
                                      quotechar=self.quotechar,
                                      data_start=self.data_start,
                                      fill_extra_cols=self.fill_extra_cols,
                                      **self.kwargs)
        conversion_info = self._read_header()
        self.check_header()
        if conversion_info is not None:
            try_int, try_float, try_string = conversion_info
        else:
            try_int = {}
            try_float = {}
            try_string = {}

        with _set_locale('C'):
            data, comments = self.engine.read(try_int, try_float, try_string)
        out = self.make_table(data, comments)

        if self.return_header_chars:
            out.meta['__ascii_fast_reader_header_chars__'] = self.engine.header_chars

        return out

    def make_table(self, data, comments):
        """Actually make the output table give the data and comments."""
        meta = OrderedDict()
        if comments:
            meta['comments'] = comments

        names = core._deduplicate_names(self.engine.get_names())
        return Table(data, names=names, meta=meta)

    def check_header(self):
        names = self.engine.get_header_names() or self.engine.get_names()
        if self.strict_names:
            # Impose strict requirements on column names (normally used in guessing)
            bads = [" ", ",", "|", "\t", "'", '"']
            for name in names:
                if (core._is_number(name)
                    or len(name) == 0
                    or name[0] in bads
                        or name[-1] in bads):
                    raise ValueError('Column name {!r} does not meet strict name requirements'
                                     .format(name))
        # When guessing require at least two columns
        if self.guessing and len(names) <= 1:
            raise ValueError('Table format guessing requires at least two columns, got {}'
                             .format(names))

    def write(self, table, output):
        """
        Use a fast Cython method to write table data to output,
        where output is a filename or file-like object.
        """
        self._write(table, output, {})

    def _write(self, table, output, default_kwargs,
               header_output=True, output_types=False):

        # Fast writer supports only 1-d columns
        core._check_multidim_table(table, max_ndim=1)

        write_kwargs = {'delimiter': self.delimiter,
                        'quotechar': self.quotechar,
                        'strip_whitespace': self.strip_whitespace_fields,
                        'comment': self.write_comment
                        }
        write_kwargs.update(default_kwargs)
        # user kwargs take precedence over default kwargs
        write_kwargs.update(self.kwargs)
        writer = cparser.FastWriter(table, **write_kwargs)
        writer.write(output, header_output, output_types)


class FastCsv(FastBasic):
    """
    A faster version of the ordinary :class:`Csv` writer that uses the
    optimized C parsing engine. Note that this reader will append empty
    field values to the end of any row with not enough columns, while
    :class:`FastBasic` simply raises an error.
    """
    _format_name = 'fast_csv'
    _description = 'Comma-separated values table using the fast C engine'
    _fast = True
    fill_extra_cols = True

    def __init__(self, **kwargs):
        super().__init__({'delimiter': ',', 'comment': None}, **kwargs)

    def write(self, table, output):
        """
        Override the default write method of `FastBasic` to
        output masked values as empty fields.
        """
        self._write(table, output, {'fill_values': [(core.masked, '')]})


class FastTab(FastBasic):
    """
    A faster version of the ordinary :class:`Tab` reader that uses
    the optimized C parsing engine.
    """
    _format_name = 'fast_tab'
    _description = 'Tab-separated values table using the fast C engine'
    _fast = True

    def __init__(self, **kwargs):
        super().__init__({'delimiter': '\t'}, **kwargs)
        self.strip_whitespace_lines = False
        self.strip_whitespace_fields = False


class FastNoHeader(FastBasic):
    """
    This class uses the fast C engine to read tables with no header line. If
    the names parameter is unspecified, the columns will be autonamed with
    "col{}".
    """
    _format_name = 'fast_no_header'
    _description = 'Basic table with no headers using the fast C engine'
    _fast = True

    def __init__(self, **kwargs):
        super().__init__({'header_start': None, 'data_start': 0}, **kwargs)

    def write(self, table, output):
        """
        Override the default writing behavior in `FastBasic` so
        that columns names are not included in output.
        """
        self._write(table, output, {}, header_output=None)


class FastCommentedHeader(FastBasic):
    """
    A faster version of the :class:`CommentedHeader` reader, which looks for
    column names in a commented line. ``header_start`` denotes the index of
    the header line among all commented lines and is 0 by default.
    """
    _format_name = 'fast_commented_header'
    _description = 'Columns name in a commented line using the fast C engine'
    _fast = True

    def __init__(self, **kwargs):
        super().__init__({}, **kwargs)
        # Mimic CommentedHeader's behavior in which data_start
        # is relative to header_start if unspecified; see #2692
        if 'data_start' not in kwargs:
            self.data_start = 0

    def make_table(self, data, comments):
        """
        Actually make the output table give the data and comments.  This is
        slightly different from the base FastBasic method in the way comments
        are handled.
        """
        meta = OrderedDict()
        if comments:
            idx = self.header_start
            if idx < 0:
                idx = len(comments) + idx
            meta['comments'] = comments[:idx] + comments[idx+1:]  # noqa
            if not meta['comments']:
                del meta['comments']

        names = core._deduplicate_names(self.engine.get_names())
        return Table(data, names=names, meta=meta)

    def _read_header(self):
        tmp = self.engine.source
        commented_lines = []

        for line in tmp.splitlines():
            line = line.lstrip()
            if line and line[0] == self.comment:  # line begins with a comment
                commented_lines.append(line[1:])
                if len(commented_lines) == self.header_start + 1:
                    break

        if len(commented_lines) <= self.header_start:
            raise cparser.CParserError('not enough commented lines')

        self.engine.setup_tokenizer([commented_lines[self.header_start]])
        self.engine.header_start = 0
        self.engine.read_header()
        self.engine.setup_tokenizer(tmp)

    def write(self, table, output):
        """
        Override the default writing behavior in `FastBasic` so
        that column names are commented.
        """
        self._write(table, output, {}, header_output='comment')


class FastRdb(FastBasic):
    """
    A faster version of the :class:`Rdb` reader. This format is similar to
    tab-delimited, but it also contains a header line after the column
    name line denoting the type of each column (N for numeric, S for string).
    """
    _format_name = 'fast_rdb'
    _description = 'Tab-separated with a type definition header line'
    _fast = True

    def __init__(self, **kwargs):
        super().__init__({'delimiter': '\t', 'data_start': 2}, **kwargs)
        self.strip_whitespace_lines = False
        self.strip_whitespace_fields = False

    def _read_header(self):
        tmp = self.engine.source
        line1 = ''
        line2 = ''
        for line in tmp.splitlines():
            # valid non-comment line
            if not line1 and line.strip() and line.lstrip()[0] != self.comment:
                line1 = line
            elif not line2 and line.strip() and line.lstrip()[0] != self.comment:
                line2 = line
                break
        else:  # less than 2 lines in table
            raise ValueError('RDB header requires 2 lines')

        # Tokenize the two header lines separately.
        # Each call to self.engine.read_header by default
        #  - calls _deduplicate_names to ensure unique header_names
        #  - sets self.names from self.header_names if not provided as kwarg
        #  - applies self.include_names/exclude_names to self.names.
        # For parsing the types disable 1+3, but self.names needs to be set.
        self.engine.setup_tokenizer([line2])
        self.engine.header_start = 0
        self.engine.read_header(deduplicate=False, filter_names=False)
        types = self.engine.get_header_names()

        # If no kwarg names have been passed, reset to have column names read from header line 1.
        if types == self.engine.get_names():
            self.engine.set_names([])
        self.engine.setup_tokenizer([line1])
        # Get full list of column names prior to applying include/exclude_names,
        # which have to be applied to the unique name set after deduplicate.
        self.engine.read_header(deduplicate=True, filter_names=False)
        col_names = self.engine.get_names()
        self.engine.read_header(deduplicate=False)
        if len(col_names) != len(types):
            raise core.InconsistentTableError('RDB header mismatch between number of '
                                              'column names and column types')
        # If columns have been removed via include/exclude_names, extract matching types.
        if len(self.engine.get_names()) != len(types):
            types = [types[col_names.index(n)] for n in self.engine.get_names()]

        if any(not re.match(r'\d*(N|S)$', x, re.IGNORECASE) for x in types):
            raise core.InconsistentTableError('RDB type definitions do not all match '
                                              '[num](N|S): {}'.format(types))

        try_int = {}
        try_float = {}
        try_string = {}

        for name, col_type in zip(self.engine.get_names(), types):
            if col_type[-1].lower() == 's':
                try_int[name] = 0
                try_float[name] = 0
                try_string[name] = 1
            else:
                try_int[name] = 1
                try_float[name] = 1
                try_string[name] = 0

        self.engine.setup_tokenizer(tmp)
        return (try_int, try_float, try_string)

    def write(self, table, output):
        """
        Override the default writing behavior in `FastBasic` to
        output a line with column types after the column name line.
        """
        self._write(table, output, {}, output_types=True)