<commit_before>from contextlib import contextmanager
import re

from alembic.testing import exclusions
from alembic.testing import TestBase, eq_, config
from alembic.testing.fixtures import op_fixture
from alembic.testing import mock
from alembic.operations import Operations
from alembic.batch import ApplyBatchImpl
from alembic.migration import MigrationContext

from sqlalchemy import Integer, Table, Column, String, MetaData, ForeignKey, \
    UniqueConstraint, ForeignKeyConstraint
from sqlalchemy.sql import column
from sqlalchemy.schema import CreateTable


class BatchApplyTest(TestBase):
    def setUp(self):
        self.op = Operations(mock.Mock(opts={}))

    def _simple_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('x', String(10)),
            Column('y', Integer)
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _uq_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('x', String()),
            Column('y', Integer),
            UniqueConstraint('y', name='uq1')
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _fk_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('email', String()),
            Column('user_id', Integer, ForeignKey('user.id'))
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _named_fk_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('email', String()),
            Column('user_id', Integer, ForeignKey('user.id', name='ufk'))
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _selfref_fk_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('parent_id', ForeignKey('tname.id')),
            Column('data', String)
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _assert_impl(self, impl, colnames=None,
                     ddl_contains=None, ddl_not_contains=None,
                     dialect='default'):
        context = op_fixture(dialect=dialect)

        impl._create(context.impl)

        if colnames is None:
            colnames = ['id', 'x', 'y']
        eq_(impl.new_table.c.keys(), colnames)

        pk_cols = [col for col in impl.new_table.c if col.primary_key]
        eq_(list(impl.new_table.primary_key), pk_cols)

        create_stmt = str(
            CreateTable(impl.new_table).compile(dialect=context.dialect))
        create_stmt = re.sub(r'[\n\t]', '', create_stmt)

        if ddl_contains:
            assert ddl_contains in create_stmt<commit_msg>- add a type here so that to account for pre-0.9 FK/type support<commit_after>from contextlib import contextmanager
import re

from alembic.testing import exclusions
from alembic.testing import TestBase, eq_, config
from alembic.testing.fixtures import op_fixture
from alembic.testing import mock
from alembic.operations import Operations
from alembic.batch import ApplyBatchImpl
from alembic.migration import MigrationContext

from sqlalchemy import Integer, Table, Column, String, MetaData, ForeignKey, \
    UniqueConstraint, ForeignKeyConstraint
from sqlalchemy.sql import column
from sqlalchemy.schema import CreateTable


class BatchApplyTest(TestBase):
    def setUp(self):
        self.op = Operations(mock.Mock(opts={}))

    def _simple_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('x', String(10)),
            Column('y', Integer)
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _uq_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('x', String()),
            Column('y', Integer),
            UniqueConstraint('y', name='uq1')
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _fk_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('email', String()),
            Column('user_id', Integer, ForeignKey('user.id'))
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _named_fk_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('email', String()),
            Column('user_id', Integer, ForeignKey('user.id', name='ufk'))
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _selfref_fk_fixture(self, table_args=(), table_kwargs={}):
        m = MetaData()
        t = Table(
            'tname', m,
            Column('id', Integer, primary_key=True),
            Column('parent_id', Integer, ForeignKey('tname.id')),
            Column('data', String)
        )
        return ApplyBatchImpl(t, table_args, table_kwargs)

    def _assert_impl(self, impl, colnames=None,
                     ddl_contains=None, ddl_not_contains=None,
                     dialect='default'):
        context = op_fixture(dialect=dialect)

        impl._create(context.impl)

        if colnames is None:
            colnames = ['id', 'x', 'y']
        eq_(impl.new_table.c.keys(), colnames)

        pk_cols = [col for col in impl.new_table.c if col.primary_key]
        eq_(list(impl.new_table.primary_key), pk_cols)

        create_stmt = str(
            CreateTable(impl.new_table).compile(dialect=context.dialect))
        create_stmt = re.sub(r'[\n\t]', '', create_stmt)

        if ddl_contains:
            assert ddl_contains in create_stmt
{'fraction_non_alphanumeric': 0.12010732323232323, 'fraction_numerical': 0.0012626262626262627, 'mean_word_length': 2.5481522956326987, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3264210', 'n_tokens_mistral': 1920, 'n_tokens_neox': 1857, 'n_words': 438}