<commit_before>from __future__ import absolute_import, division, print_function

import os
import re
import subprocess

from itertools import chain
from collections import Iterator
from datetime import datetime, date
from distutils.spawn import find_executable

import pandas as pd
import sqlalchemy as sa
from sqlalchemy import inspect
from sqlalchemy.ext.compiler import compiles
from sqlalchemy import event
from sqlalchemy.schema import CreateSchema

from multipledispatch import MDNotImplementedError

import datashape
from datashape import DataShape, Record, Option, var, dshape
from datashape.predicates import isdimension, isrecord, isscalar
from datashape import discover
from datashape.dispatch import dispatch

from toolz import (partition_all, keyfilter, first, memoize, valfilter,
                   identity, concat, curry, merge)
from toolz.curried import pluck, map

from ..compatibility import unicode
from ..utils import keywords, ignoring, iter_except
from ..convert import convert, ooc_types
from ..append import append
from ..resource import resource
from ..chunks import Chunks
from .csv import CSV

base = (int, float, datetime, date, bool, str)


# http://docs.sqlalchemy.org/en/latest/core/types.html

types = {
    'int64': sa.types.BigInteger,
    'int32': sa.types.Integer,
    'int': sa.types.Integer,
    'int16': sa.types.SmallInteger,
    'float32': sa.types.Float(precision=24),  # sqlalchemy uses mantissa
    'float64': sa.types.Float(precision=53),  # for precision
    'float': sa.types.Float(precision=53),
    'real': sa.types.Float(precision=53),
    'string': sa.types.Text,
    'date': sa.types.Date,
    'time': sa.types.Time,
    'datetime': sa.types.DateTime,
    'bool': sa.types.Boolean,
    "timedelta[unit='D']": sa.types.Interval(second_precision=0,
                                             day_precision=9),
    "timedelta[unit='h']": sa.types.Interval(second_precision=0,
                                             day_precision=0),
    "timedelta[unit='m']": sa.types.Interval(second_precision=0,
                                             day_precision=0),
    "timedelta[unit='s']": sa.types.Interval(second_precision=0,
                                             day_precision=0),
    "timedelta[unit='ms']": sa.types.Interval(second_precision=3,
                                              day_precision=0),
    "timedelta[unit='us']": sa.types.Interval(second_precision=6,
                                              day_precision=0),
    "timedelta[unit='ns']": sa.types.Interval(second_precision=9,
                                              day_precision=0),
    # ??: sa.types.LargeBinary,
    # Decimal: sa.types.Numeric,
    # ??: sa.types.PickleType,
    # unicode: sa.types.Unicode,
    # unicode: sa.types.UnicodeText,
    # str: sa.types.Text,  # ??
}

revtypes = dict(map(reversed, types.items()))

revtypes.update({
    sa.types.VARCHAR: 'string',
    sa.types.String: 'string',
    sa.types.Unicode: 'string',
    sa.types.DATETIME: 'datetime',
    sa.types.TIMESTAMP: 'datetime',
    sa.types.FLOAT: 'float64',
    sa.types.DATE: 'date',
    sa.types.BIGINT: 'int64',
    sa.types.INTEGER: 'int',
    sa.types.NUMERIC: 'float64',  # TODO: extend datashape to decimal
    sa.types.BIGINT: 'int64',
    sa.types.NullType: 'string',
    sa.types.Float: 'float64',
})

# interval types are special cased in discover_typeengine so remove them from
# revtypes
revtypes = valfilter(lambda x: not isinstance(x, sa.types.Interval), revtypes)


units_of_power = {
    0: 's',
    3: 'ms',
    6: 'us',
    9: 'ns'
}

# these aren't loaded by sqlalchemy by default
sa.dialects.registry.load('oracle')
sa.dialects.registry.load('postgresql')


def batch(sel, chunksize=10000):
    """Execute `sel`, streaming row at a time and fetching from the database in
    batches of size `chunksize`.

    Parameters
    ----------
    sel : sa.sql.Selectable
        Selectable to execute
    chunksize : int, optional, default 10000
        Number of rows to fetch from the database
    """
    def rowterator(sel, chunksize):
        with sel.bind.begin() as conn:
            result = conn.execution_options(stream_results=True).execute(sel)
            yield result.keys()

            for rows in iter_except(curry(result.fetchmany, size=chunksize),
                                    sa.exc.ResourceClosedError):
                if rows:
                    yield rows
                else:
                    return
    terator = rowterator(sel, chunksize)
    return next(terator), concat(terator)


@discover.register(sa.sql.elements.ColumnClause)
def discover_column_clause(t):
    return Record([(t.name, discover(t.type))])


@discover.register(sa.dialects.postgresql.base.INTERVAL)
def discover_postgresql_interval(t):
    return discover(sa.types.Interval(day_precision=0,
                                      second_precision=t.precision))


@discover.register(sa.dialects.oracle.base.INTERVAL)
def discover_oracle_interval(t):
    return discover(t.adapt(sa.types.Interval))


@discover.register(sa.sql.type_api.TypeEngine)
def discover_typeengine(typ):
    if isinstance(typ, sa.types.Interval):
        if typ.second_precision is None and typ.day_precision is None:
            return datashape.TimeDelta(unit='us')
        elif typ.second_precision == 0 and typ.day_precision == 0:
            return datashape.TimeDelta(unit='s')

        if typ.second_precision in units_of_power and not typ.day_precision:
            units = units_of_power[typ.second_precision]
        elif typ.day_precision > 0:
            units = 'D'
        else:
            raise ValueError('Cannot infer INTERVAL type with parameters'
                             'second_precision=%d, day_precision=%d' %
                             (typ.second_precision, typ.day_precision))
        return datashape.TimeDelta(unit=units)
    if typ in revtypes:
        return dshape(revtypes[typ])[0]
    if type(typ) in revtypes:
        return dshape(revtypes[type(typ)])[0]
    else:
        for k, v in revtypes.items():
            if isinstance(k, type) and (isinstance(typ, k) or
                                        hasattr(typ, 'impl') and
                                        isinstance(typ.impl, k)):
                return v
            if k == typ:
                return v
    raise NotImplementedError("No SQL-datashape match for type %s" % typ)


@discover.register(sa.Column)
def discover_sqlalchemy_column(col):
    optionify = Option if col.nullable else identity
    return Record([[col.name, optionify(discover(col.type))]])


@discover.register(sa.sql.FromClause)
def discover_sqlalchemy_selectable(t):
    records = list(sum([discover(c).parameters[0] for c in t.columns], ()))
    return var * Record(records)


@memoize
def metadata_of_engine(engine, schema=None):
    return sa.MetaData(engine, schema=schema)


def create_engine(uri, *args, **kwargs):<commit_msg>Make this a little easier to debug<commit_after>from __future__ import absolute_import, division, print_function

import os
import re
import subprocess

from itertools import chain
from collections import Iterator
from datetime import datetime, date
from distutils.spawn import find_executable

import pandas as pd
import sqlalchemy as sa
from sqlalchemy import inspect
from sqlalchemy.ext.compiler import compiles
from sqlalchemy import event
from sqlalchemy.schema import CreateSchema

from multipledispatch import MDNotImplementedError

import datashape
from datashape import DataShape, Record, Option, var, dshape
from datashape.predicates import isdimension, isrecord, isscalar
from datashape import discover
from datashape.dispatch import dispatch

from toolz import (partition_all, keyfilter, first, memoize, valfilter,
                   identity, concat, curry, merge)
from toolz.curried import pluck, map

from ..compatibility import unicode
from ..utils import keywords, ignoring, iter_except
from ..convert import convert, ooc_types
from ..append import append
from ..resource import resource
from ..chunks import Chunks
from .csv import CSV

base = (int, float, datetime, date, bool, str)


# http://docs.sqlalchemy.org/en/latest/core/types.html

types = {
    'int64': sa.types.BigInteger,
    'int32': sa.types.Integer,
    'int': sa.types.Integer,
    'int16': sa.types.SmallInteger,
    'float32': sa.types.Float(precision=24),  # sqlalchemy uses mantissa
    'float64': sa.types.Float(precision=53),  # for precision
    'float': sa.types.Float(precision=53),
    'real': sa.types.Float(precision=53),
    'string': sa.types.Text,
    'date': sa.types.Date,
    'time': sa.types.Time,
    'datetime': sa.types.DateTime,
    'bool': sa.types.Boolean,
    "timedelta[unit='D']": sa.types.Interval(second_precision=0,
                                             day_precision=9),
    "timedelta[unit='h']": sa.types.Interval(second_precision=0,
                                             day_precision=0),
    "timedelta[unit='m']": sa.types.Interval(second_precision=0,
                                             day_precision=0),
    "timedelta[unit='s']": sa.types.Interval(second_precision=0,
                                             day_precision=0),
    "timedelta[unit='ms']": sa.types.Interval(second_precision=3,
                                              day_precision=0),
    "timedelta[unit='us']": sa.types.Interval(second_precision=6,
                                              day_precision=0),
    "timedelta[unit='ns']": sa.types.Interval(second_precision=9,
                                              day_precision=0),
    # ??: sa.types.LargeBinary,
    # Decimal: sa.types.Numeric,
    # ??: sa.types.PickleType,
    # unicode: sa.types.Unicode,
    # unicode: sa.types.UnicodeText,
    # str: sa.types.Text,  # ??
}

revtypes = dict(map(reversed, types.items()))

revtypes.update({
    sa.types.VARCHAR: 'string',
    sa.types.String: 'string',
    sa.types.Unicode: 'string',
    sa.types.DATETIME: 'datetime',
    sa.types.TIMESTAMP: 'datetime',
    sa.types.FLOAT: 'float64',
    sa.types.DATE: 'date',
    sa.types.BIGINT: 'int64',
    sa.types.INTEGER: 'int',
    sa.types.NUMERIC: 'float64',  # TODO: extend datashape to decimal
    sa.types.BIGINT: 'int64',
    sa.types.NullType: 'string',
    sa.types.Float: 'float64',
})

# interval types are special cased in discover_typeengine so remove them from
# revtypes
revtypes = valfilter(lambda x: not isinstance(x, sa.types.Interval), revtypes)


units_of_power = {
    0: 's',
    3: 'ms',
    6: 'us',
    9: 'ns'
}

# these aren't loaded by sqlalchemy by default
sa.dialects.registry.load('oracle')
sa.dialects.registry.load('postgresql')


def batch(sel, chunksize=10000):
    """Execute `sel`, streaming row at a time and fetching from the database in
    batches of size `chunksize`.

    Parameters
    ----------
    sel : sa.sql.Selectable
        Selectable to execute
    chunksize : int, optional, default 10000
        Number of rows to fetch from the database
    """
    def rowterator(sel, chunksize):
        with sel.bind.begin() as conn:
            result = conn.execution_options(stream_results=True).execute(sel)
            yield result.keys()

            for rows in iter_except(curry(result.fetchmany, size=chunksize),
                                    sa.exc.ResourceClosedError):
                if rows:
                    yield rows
                else:
                    return
    terator = rowterator(sel, chunksize)
    return next(terator), concat(terator)


@discover.register(sa.sql.elements.ColumnClause)
def discover_column_clause(t):
    return Record([(t.name, discover(t.type))])


@discover.register(sa.dialects.postgresql.base.INTERVAL)
def discover_postgresql_interval(t):
    return discover(sa.types.Interval(day_precision=0,
                                      second_precision=t.precision))


@discover.register(sa.dialects.oracle.base.INTERVAL)
def discover_oracle_interval(t):
    return discover(t.adapt(sa.types.Interval))


@discover.register(sa.sql.type_api.TypeEngine)
def discover_typeengine(typ):
    if isinstance(typ, sa.types.Interval):
        if typ.second_precision is None and typ.day_precision is None:
            return datashape.TimeDelta(unit='us')
        elif typ.second_precision == 0 and typ.day_precision == 0:
            return datashape.TimeDelta(unit='s')

        if typ.second_precision in units_of_power and not typ.day_precision:
            units = units_of_power[typ.second_precision]
        elif typ.day_precision > 0:
            units = 'D'
        else:
            raise ValueError('Cannot infer INTERVAL type with parameters'
                             'second_precision=%d, day_precision=%d' %
                             (typ.second_precision, typ.day_precision))
        return datashape.TimeDelta(unit=units)
    if typ in revtypes:
        return dshape(revtypes[typ])[0]
    if type(typ) in revtypes:
        return dshape(revtypes[type(typ)])[0]
    else:
        for k, v in revtypes.items():
            if isinstance(k, type) and (isinstance(typ, k) or
                                        hasattr(typ, 'impl') and
                                        isinstance(typ.impl, k)):
                return v
            if k == typ:
                return v
    raise NotImplementedError("No SQL-datashape match for type %s" % typ)


@discover.register(sa.Column)
def discover_sqlalchemy_column(col):
    optionify = Option if col.nullable else identity
    return Record([[col.name, optionify(discover(col.type))]])


@discover.register(sa.sql.FromClause)
def discover_sqlalchemy_selectable(t):
    types = [discover(c).parameters[0] for c in t.columns]
    return var * Record(list(sum(types, ())))


@memoize
def metadata_of_engine(engine, schema=None):
    return sa.MetaData(engine, schema=schema)


def create_engine(uri, *args, **kwargs):
{'fraction_non_alphanumeric': 0.11191724732418648, 'fraction_numerical': 0.009051073917103657, 'mean_word_length': 2.655027566290365, 'pattern_counts': {'":': 14, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 36, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5353202', 'n_tokens_mistral': 4294, 'n_tokens_neox': 4097, 'n_words': 983}