CLN: Consistent imports in tests: pandas.util.testing as tm

[SaturnFromTitan]

This is a follow-up issue of #23018 with the goal of enforcing a consistent approach of importing things from pandas.util.testing inside the test suite.

Currently, there are multiple import styles used:

1. import pandas.util.testing as tm
2. from pandas.util.testing import ...
3. from pandas.util import testing as tm

As it is already the dominant import style used and also less prone to circular imports 1. is the prefered option. After migrating all current occurrences of 2. and 3. we'll add a linting rule that enforces this convention.

These are the current files that use 2. (generated with grep -R -l --include="*.py" "from pandas.util.testing import " pandas/tests:

• pandas/tests/arrays/sparse/test_array.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/computation/test_eval.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/frame/test_api.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part1 #29300
• pandas/tests/frame/test_apply.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part1 #29300
• pandas/tests/frame/test_axis_select_reindex.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part1 #29300
• pandas/tests/frame/test_block_internals.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part1 #29300
• pandas/tests/frame/test_combine_concat.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part1 #29300
• pandas/tests/frame/test_dtypes.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part1 #29300
• pandas/tests/frame/test_indexing.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 2 #29301
• pandas/tests/frame/test_missing.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 2 #29301
• pandas/tests/frame/test_mutate_columns.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 2 #29301
• pandas/tests/frame/test_nonunique_indexes.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 2 #29301
• pandas/tests/frame/test_operators.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 2 #29301
• pandas/tests/frame/test_quantile.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 2 #29301
• pandas/tests/frame/test_query_eval.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 3 #29302
• pandas/tests/frame/test_rank.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 3 #29302
• pandas/tests/frame/test_replace.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 3 #29302
• pandas/tests/frame/test_reshape.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 3 #29302
• pandas/tests/frame/test_sort_values_level_as_str.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 4 #29303
• pandas/tests/frame/test_sorting.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 4 #29303
• pandas/tests/frame/test_timeseries.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 4 #29303
• pandas/tests/frame/test_to_csv.py CLN: Consistent pandas.util.testing imports in pandas/tests/frame - part 4 #29303
• pandas/tests/generic/test_frame.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/generic/test_generic.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/generic/test_series.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/groupby/test_categorical.py CLN: Consistent pandas.util.testing imports in pandas/tests/groupby #29287
• pandas/tests/groupby/test_counting.py CLN: Consistent pandas.util.testing imports in pandas/tests/groupby #29287
• pandas/tests/groupby/test_groupby.py CLN: Consistent pandas.util.testing imports in pandas/tests/groupby #29287
• pandas/tests/groupby/test_grouping.py CLN: Consistent pandas.util.testing imports in pandas/tests/groupby #29287
• pandas/tests/groupby/test_index_as_string.py CLN: Consistent pandas.util.testing imports in pandas/tests/groupby #29287
• pandas/tests/groupby/test_nth.py CLN: Consistent pandas.util.testing imports in pandas/tests/groupby #29287
• pandas/tests/groupby/test_timegrouper.py CLN: Consistent pandas.util.testing imports in pandas/tests/groupby #29287
• pandas/tests/groupby/test_transform.py CLN: Consistent pandas.util.testing imports in pandas/tests/groupby #29287
• pandas/tests/indexes/datetimes/test_datetime.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexes/datetimes/test_tools.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexes/multi/test_astype.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexes/multi/test_indexing.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexes/test_base.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexes/test_category.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexes/timedeltas/test_partial_slicing.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexes/timedeltas/test_timedelta.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexes/timedeltas/test_tools.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexes #29289
• pandas/tests/indexing/test_categorical.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexing #29291
• pandas/tests/indexing/test_floats.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexing #29291
• pandas/tests/indexing/test_iloc.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexing #29291
• pandas/tests/indexing/test_indexing.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexing #29291
• pandas/tests/indexing/test_ix.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexing #29291
• pandas/tests/indexing/test_loc.py CLN: Consistent pandas.util.testing imports in pandas/tests/indexing #29291
• pandas/tests/internals/test_internals.py
• pandas/tests/io/excel/test_openpyxl.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/excel/test_style.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/excel/test_writers.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/excel/test_xlrd.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/excel/test_xlsxwriter.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/excel/test_xlwt.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/json/test_compression.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/json/test_pandas.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/json/test_readlines.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/parser/test_textreader.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/pytables/test_compat.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/pytables/test_complex.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/pytables/test_store.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/pytables/test_timezones.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/test_clipboard.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/test_feather.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/test_gcs.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/test_html.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/io/test_packers.py CLN: Consistent pandas.util.testing imports in pandas/tests/io #29299
• pandas/tests/plotting/common.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/plotting/test_datetimelike.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/resample/test_base.py CLN: Consistent pandas.util.testing imports in pandas/tests/resample #29286
• pandas/tests/resample/test_datetime_index.py CLN: Consistent pandas.util.testing imports in pandas/tests/resample #29286
• pandas/tests/resample/test_period_index.py CLN: Consistent pandas.util.testing imports in pandas/tests/resample #29286
• pandas/tests/resample/test_resample_api.py CLN: Consistent pandas.util.testing imports in pandas/tests/resample #29286
• pandas/tests/resample/test_resampler_grouper.py CLN: Consistent pandas.util.testing imports in pandas/tests/resample #29286
• pandas/tests/resample/test_time_grouper.py CLN: Consistent pandas.util.testing imports in pandas/tests/resample #29286
• pandas/tests/resample/test_timedelta.py CLN: Consistent pandas.util.testing imports in pandas/tests/resample #29286
• pandas/tests/reshape/merge/test_join.py CLN: Consistent pandas.util.testing imports in pandas/tests/reshape #29285
• pandas/tests/reshape/merge/test_merge.py CLN: Consistent pandas.util.testing imports in pandas/tests/reshape #29285
• pandas/tests/reshape/merge/test_merge_asof.py CLN: Consistent pandas.util.testing imports in pandas/tests/reshape #29285
• pandas/tests/reshape/merge/test_merge_index_as_string.py CLN: Consistent pandas.util.testing imports in pandas/tests/reshape #29285
• pandas/tests/reshape/merge/test_merge_ordered.py CLN: Consistent pandas.util.testing imports in pandas/tests/reshape #29285
• pandas/tests/reshape/test_concat.py CLN: Consistent pandas.util.testing imports in pandas/tests/reshape #29285
• pandas/tests/reshape/test_reshape.py CLN: Consistent pandas.util.testing imports in pandas/tests/reshape #29285
• pandas/tests/series/indexing/test_alter_index.py CLN: Consistent pandas.util.testing imports in pandas/tests/series/indexing #29277
• pandas/tests/series/indexing/test_boolean.py CLN: Consistent pandas.util.testing imports in pandas/tests/series/indexing #29277
• pandas/tests/series/indexing/test_datetime.py CLN: Consistent pandas.util.testing imports in pandas/tests/series/indexing #29277
• pandas/tests/series/indexing/test_iloc.py CLN: Consistent pandas.util.testing imports in pandas/tests/series/indexing #29277
• pandas/tests/series/indexing/test_indexing.py CLN: Consistent pandas.util.testing imports in pandas/tests/series/indexing #29277
• pandas/tests/series/indexing/test_loc.py CLN: Consistent pandas.util.testing imports in pandas/tests/series/indexing #29277
• pandas/tests/series/indexing/test_numeric.py CLN: Consistent pandas.util.testing imports in pandas/tests/series/indexing #29277
• pandas/tests/series/test_analytics.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_api.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_apply.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_combine_concat.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_constructors.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_datetime_values.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_internals.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_io.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_missing.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_operators.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_rank.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_sorting.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/series/test_timeseries.py CLN: Consistent pandas.util.testing imports in pandas/tests/series #29273
• pandas/tests/test_algos.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/test_expressions.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/test_join.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/test_sorting.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/test_strings.py CLN: Consistent pandas.util.testing imports in remaining test suite #29307
• pandas/tests/util/test_assert_almost_equal.py CLN: Consistent pandas.util.testing imports in pandas/tests/util #29278
• pandas/tests/util/test_assert_categorical_equal.py CLN: Consistent pandas.util.testing imports in pandas/tests/util #29278
• pandas/tests/util/test_assert_extension_array_equal.py CLN: Consistent pandas.util.testing imports in pandas/tests/util #29278
• pandas/tests/util/test_assert_frame_equal.py CLN: Consistent pandas.util.testing imports in pandas/tests/util #29278
• pandas/tests/util/test_assert_index_equal.py CLN: Consistent pandas.util.testing imports in pandas/tests/util #29278
• pandas/tests/util/test_assert_interval_array_equal.py CLN: Consistent pandas.util.testing imports in pandas/tests/util #29278
• pandas/tests/util/test_assert_numpy_array_equal.py CLN: Consistent pandas.util.testing imports in pandas/tests/util #29278
• pandas/tests/util/test_assert_series_equal.py CLN: Consistent pandas.util.testing imports in pandas/tests/util #29278

[SaturnFromTitan]

After investigating a few files manually, I think this should be (mostly) scriptable:

1. Understand which items are imported directly from pandas.util.testing. Then do the following for all of them (for x in ...):
   1. Replace x with tm.x in the whole file
   2. Replace tm.tm.x with tm.x in the whole file (needed because the step before broke some code)
2. Afterwards, remove the (now broken) imports
3. If necessary, add import pandas.util.testing as tm

[SaturnFromTitan]

I came up with the following script which should solve the problem pretty well (+should be rather straight forward to adjust to other use cases):

import ast
import copy
import pathlib

TM_MODULE = "pandas.util.testing"
FROM_TM_IMPORT_STATEMENT = f"from {TM_MODULE} import"
TM_IMPORT_STATEMENT = f"import {TM_MODULE} as tm"


def read_file(path):
    with open(path, "r") as f:
        print(f"\nReading {path}...")
        return f.read()


def get_ast(content):
    return ast.parse(content)


def has_from_tm_imports(content):
    return FROM_TM_IMPORT_STATEMENT in content


def add_tm_import_if_needed(content):
    if TM_IMPORT_STATEMENT in content:
        # statement is already available
        return content

    if "import" not in content:
        # if nothing is imported, then there's nothing to fix
        return content

    lines = content.split("\n")
    new_lines = copy.deepcopy(lines)
    for i, line in enumerate(lines):
        if "import" in line:
            # add TM_IMPORT_STATEMENT as first import.
            # isort will later resolve the order
            new_lines.insert(i, TM_IMPORT_STATEMENT)
            break
    return "\n".join(new_lines)


def get_direct_imports_to_replace(tree):
    for node in tree.body:
        if isinstance(node, ast.ImportFrom):
            if node.module == TM_MODULE:
                return [n.name for n in node.names]
    raise Exception(f"Couldn't find any direct imports from {TM_MODULE}")


def determine_direct_tm_import_lines(tree):
    start_line = -1
    end_line = -1

    last_was_tm_import = False
    for node in tree.body:
        if last_was_tm_import:
            end_line = node.lineno
            last_was_tm_import = False

        if isinstance(node, ast.ImportFrom) and node.module == TM_MODULE:
            start_line = node.lineno
            last_was_tm_import = True

    assert start_line != -1, "Can't determine start line of from tm import"
    assert end_line != -1, "Can't determine end line of from tm import"

    return start_line, end_line


def remove_direct_tm_import_statements(content):
    """ Get the line number of the unwanted import and the line number of the next statement.

    Everything in between will be removed. This might yield too many empty lines
    removed, but black can fix that later on.
    """
    start_line, end_line = determine_direct_tm_import_lines(tree)
    start_idx = start_line - 1
    end_idx = end_line - 1

    lines = content.split("\n")
    end_idx = check_end_idx(lines, start_idx, end_idx)

    new_lines = lines[:start_idx] + lines[end_idx:]
    return "\n".join(new_lines)


def check_end_idx(lines, start_idx, end_idx):
    """
    end_idx is by default the index of the next statement. This may result in the
    deletion of comments or empty lines which isn't intended. In this case we
    can be "sharper" with end_idx.
    """
    assert start_idx < end_idx, "start_idx must be smaller than end_idx"
    for idx in range(start_idx + 1, end_idx):
        line = lines[idx]
        is_empty_line = line == ""
        is_single_comment = line.startswith("#")
        is_multi_comment = line.startswith('"""') or line.startswith("'''")
        if any([is_empty_line, is_single_comment, is_multi_comment]):
            return idx
    return end_idx


def replace_direct_imports(content, direct_imports):
    """
    Say we have a direct import x. Then in the given file it might be imported as
    'x' or 'tm.x'

    So first we replace 'x' with 'tm.x' which yields
        x -> tm.x
        tm.x ->tm.tm.x

    Now we have to take care of the 'tm.tm.x' issues, so we simply replace 'tm.tm.x' with 'x'
        tm.x -> tm.x
        tm.tm.x -> tm.x

    Aaaand we're done!

    EDIT: Nearly...
    Some modules use helpers like _assert_series_equal which will be transformed to _tm.assert_series_equal which is
    not a valid function name anymore. Therefore we have to add one final transformation '_.tm' to '_'
    """
    for name in direct_imports:
        content = content.replace(name, f"tm.{name}")

    for name in direct_imports:
        content = content.replace(f"tm.tm.{name}", f"tm.{name}")

    content = content.replace("_tm.", "_")
    return content


def overwrite_file(path, content):
    with open(path, "w") as f:
        print(f"Writing {path}...")
        f.write(content)


if __name__ == "__main__":
    directory = "pandas/tests"

    pathlist = pathlib.Path(directory).glob("**/*.py")
    for path in pathlist:
        path_str = str(path)

        content = read_file(path_str)
        if not has_from_tm_imports(content):
            continue

        content = add_tm_import_if_needed(content)
        tree = get_ast(content)
        direct_imports = get_direct_imports_to_replace(tree)
        content = remove_direct_tm_import_statements(content)
        content = replace_direct_imports(content, direct_imports)

        overwrite_file(path_str, content)

I'll open a few PRs for the remaining files :)