from __future__ import annotations import collections from datetime import datetime from decimal import Decimal from functools import wraps import operator import os import re import string from typing import ( TYPE_CHECKING, Callable, ContextManager, Counter, Iterable, ) import warnings import numpy as np from pandas._config.localization import ( # noqa:F401 can_set_locale, get_locales, set_locale, ) from pandas._typing import Dtype from pandas.core.dtypes.common import ( is_float_dtype, is_integer_dtype, is_sequence, is_unsigned_integer_dtype, pandas_dtype, ) import pandas as pd from pandas import ( Categorical, CategoricalIndex, DataFrame, DatetimeIndex, Index, IntervalIndex, MultiIndex, RangeIndex, Series, bdate_range, ) from pandas._testing._io import ( # noqa:F401 close, network, round_trip_localpath, round_trip_pathlib, round_trip_pickle, with_connectivity_check, write_to_compressed, ) from pandas._testing._random import ( # noqa:F401 randbool, rands, rands_array, randu_array, ) from pandas._testing._warnings import assert_produces_warning # noqa:F401 from pandas._testing.asserters import ( # noqa:F401 assert_almost_equal, assert_attr_equal, assert_categorical_equal, assert_class_equal, assert_contains_all, assert_copy, assert_datetime_array_equal, assert_dict_equal, assert_equal, assert_extension_array_equal, assert_frame_equal, assert_index_equal, assert_indexing_slices_equivalent, assert_interval_array_equal, assert_is_sorted, assert_is_valid_plot_return_object, assert_metadata_equivalent, assert_numpy_array_equal, assert_period_array_equal, assert_series_equal, assert_sp_array_equal, assert_timedelta_array_equal, raise_assert_detail, ) from pandas._testing.compat import ( # noqa:F401 get_dtype, get_obj, ) from pandas._testing.contexts import ( # noqa:F401 RNGContext, decompress_file, ensure_clean, ensure_clean_dir, ensure_safe_environment_variables, set_timezone, use_numexpr, with_csv_dialect, ) from pandas.core.api import ( Float64Index, Int64Index, NumericIndex, UInt64Index, ) from pandas.core.arrays import ( BaseMaskedArray, ExtensionArray, PandasArray, ) from pandas.core.arrays._mixins import NDArrayBackedExtensionArray from pandas.core.construction import extract_array if TYPE_CHECKING: from pandas import ( PeriodIndex, TimedeltaIndex, ) _N = 30 _K = 4 UNSIGNED_INT_NUMPY_DTYPES: list[Dtype] = ["uint8", "uint16", "uint32", "uint64"] UNSIGNED_INT_EA_DTYPES: list[Dtype] = ["UInt8", "UInt16", "UInt32", "UInt64"] SIGNED_INT_NUMPY_DTYPES: list[Dtype] = [int, "int8", "int16", "int32", "int64"] SIGNED_INT_EA_DTYPES: list[Dtype] = ["Int8", "Int16", "Int32", "Int64"] ALL_INT_NUMPY_DTYPES = UNSIGNED_INT_NUMPY_DTYPES + SIGNED_INT_NUMPY_DTYPES ALL_INT_EA_DTYPES = UNSIGNED_INT_EA_DTYPES + SIGNED_INT_EA_DTYPES FLOAT_NUMPY_DTYPES: list[Dtype] = [float, "float32", "float64"] FLOAT_EA_DTYPES: list[Dtype] = ["Float32", "Float64"] COMPLEX_DTYPES: list[Dtype] = [complex, "complex64", "complex128"] STRING_DTYPES: list[Dtype] = [str, "str", "U"] DATETIME64_DTYPES: list[Dtype] = ["datetime64[ns]", "M8[ns]"] TIMEDELTA64_DTYPES: list[Dtype] = ["timedelta64[ns]", "m8[ns]"] BOOL_DTYPES: list[Dtype] = [bool, "bool"] BYTES_DTYPES: list[Dtype] = [bytes, "bytes"] OBJECT_DTYPES: list[Dtype] = [object, "object"] ALL_REAL_NUMPY_DTYPES = FLOAT_NUMPY_DTYPES + ALL_INT_NUMPY_DTYPES ALL_NUMPY_DTYPES = ( ALL_REAL_NUMPY_DTYPES + COMPLEX_DTYPES + STRING_DTYPES + DATETIME64_DTYPES + TIMEDELTA64_DTYPES + BOOL_DTYPES + OBJECT_DTYPES + BYTES_DTYPES ) NARROW_NP_DTYPES = [ np.float16, np.float32, np.int8, np.int16, np.int32, np.uint8, np.uint16, np.uint32, ] NULL_OBJECTS = [None, np.nan, pd.NaT, float("nan"), pd.NA, Decimal("NaN")] NP_NAT_OBJECTS = [ cls("NaT", unit) for cls in [np.datetime64, np.timedelta64] for unit in [ "Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns", "ps", "fs", "as", ] ] EMPTY_STRING_PATTERN = re.compile("^$") # set testing_mode _testing_mode_warnings = (DeprecationWarning, ResourceWarning) def set_testing_mode(): # set the testing mode filters testing_mode = os.environ.get("PANDAS_TESTING_MODE", "None") if "deprecate" in testing_mode: for category in _testing_mode_warnings: warnings.simplefilter("always", category) def reset_testing_mode(): # reset the testing mode filters testing_mode = os.environ.get("PANDAS_TESTING_MODE", "None") if "deprecate" in testing_mode: for category in _testing_mode_warnings: warnings.simplefilter("ignore", category) set_testing_mode() def reset_display_options(): """ Reset the display options for printing and representing objects. """ pd.reset_option("^display.", silent=True) # ----------------------------------------------------------------------------- # Comparators def equalContents(arr1, arr2) -> bool: """ Checks if the set of unique elements of arr1 and arr2 are equivalent. """ return frozenset(arr1) == frozenset(arr2) def box_expected(expected, box_cls, transpose=True): """ Helper function to wrap the expected output of a test in a given box_class. Parameters ---------- expected : np.ndarray, Index, Series box_cls : {Index, Series, DataFrame} Returns ------- subclass of box_cls """ if box_cls is pd.array: if isinstance(expected, RangeIndex): # pd.array would return an IntegerArray expected = PandasArray(np.asarray(expected._values)) else: expected = pd.array(expected) elif box_cls is Index: expected = Index._with_infer(expected) elif box_cls is Series: expected = Series(expected) elif box_cls is DataFrame: expected = Series(expected).to_frame() if transpose: # for vector operations, we need a DataFrame to be a single-row, # not a single-column, in order to operate against non-DataFrame # vectors of the same length. But convert to two rows to avoid # single-row special cases in datetime arithmetic expected = expected.T expected = pd.concat([expected] * 2, ignore_index=True) elif box_cls is np.ndarray or box_cls is np.array: expected = np.array(expected) elif box_cls is to_array: expected = to_array(expected) else: raise NotImplementedError(box_cls) return expected def to_array(obj): """ Similar to pd.array, but does not cast numpy dtypes to nullable dtypes. """ # temporary implementation until we get pd.array in place dtype = getattr(obj, "dtype", None) if dtype is None: return np.asarray(obj) return extract_array(obj, extract_numpy=True) # ----------------------------------------------------------------------------- # Others def getCols(k): return string.ascii_uppercase[:k] # make index def makeStringIndex(k=10, name=None): return Index(rands_array(nchars=10, size=k), name=name) def makeUnicodeIndex(k=10, name=None): return Index(randu_array(nchars=10, size=k), name=name) def makeCategoricalIndex(k=10, n=3, name=None, **kwargs): """make a length k index or n categories""" x = rands_array(nchars=4, size=n) return CategoricalIndex( Categorical.from_codes(np.arange(k) % n, categories=x), name=name, **kwargs ) def makeIntervalIndex(k=10, name=None, **kwargs): """make a length k IntervalIndex""" x = np.linspace(0, 100, num=(k + 1)) return IntervalIndex.from_breaks(x, name=name, **kwargs) def makeBoolIndex(k=10, name=None): if k == 1: return Index([True], name=name) elif k == 2: return Index([False, True], name=name) return Index([False, True] + [False] * (k - 2), name=name) def makeNumericIndex(k=10, name=None, *, dtype): dtype = pandas_dtype(dtype) assert isinstance(dtype, np.dtype) if is_integer_dtype(dtype): values = np.arange(k, dtype=dtype) if is_unsigned_integer_dtype(dtype): values += 2 ** (dtype.itemsize * 8 - 1) elif is_float_dtype(dtype): values = np.random.random_sample(k) - np.random.random_sample(1) values.sort() values = values * (10 ** np.random.randint(0, 9)) else: raise NotImplementedError(f"wrong dtype {dtype}") return NumericIndex(values, dtype=dtype, name=name) def makeIntIndex(k=10, name=None): base_idx = makeNumericIndex(k, name=name, dtype="int64") return Int64Index(base_idx) def makeUIntIndex(k=10, name=None): base_idx = makeNumericIndex(k, name=name, dtype="uint64") return UInt64Index(base_idx) def makeRangeIndex(k=10, name=None, **kwargs): return RangeIndex(0, k, 1, name=name, **kwargs) def makeFloatIndex(k=10, name=None): base_idx = makeNumericIndex(k, name=name, dtype="float64") return Float64Index(base_idx) def makeDateIndex(k: int = 10, freq="B", name=None, **kwargs) -> DatetimeIndex: dt = datetime(2000, 1, 1) dr = bdate_range(dt, periods=k, freq=freq, name=name) return DatetimeIndex(dr, name=name, **kwargs) def makeTimedeltaIndex(k: int = 10, freq="D", name=None, **kwargs) -> TimedeltaIndex: return pd.timedelta_range(start="1 day", periods=k, freq=freq, name=name, **kwargs) def makePeriodIndex(k: int = 10, name=None, **kwargs) -> PeriodIndex: dt = datetime(2000, 1, 1) return pd.period_range(start=dt, periods=k, freq="B", name=name, **kwargs) def makeMultiIndex(k=10, names=None, **kwargs): return MultiIndex.from_product((("foo", "bar"), (1, 2)), names=names, **kwargs) _names = [ "Alice", "Bob", "Charlie", "Dan", "Edith", "Frank", "George", "Hannah", "Ingrid", "Jerry", "Kevin", "Laura", "Michael", "Norbert", "Oliver", "Patricia", "Quinn", "Ray", "Sarah", "Tim", "Ursula", "Victor", "Wendy", "Xavier", "Yvonne", "Zelda", ] def _make_timeseries(start="2000-01-01", end="2000-12-31", freq="1D", seed=None): """ Make a DataFrame with a DatetimeIndex Parameters ---------- start : str or Timestamp, default "2000-01-01" The start of the index. Passed to date_range with `freq`. end : str or Timestamp, default "2000-12-31" The end of the index. Passed to date_range with `freq`. freq : str or Freq The frequency to use for the DatetimeIndex seed : int, optional The random state seed. * name : object dtype with string names * id : int dtype with * x, y : float dtype Examples -------- >>> _make_timeseries() # doctest: +SKIP id name x y timestamp 2000-01-01 982 Frank 0.031261 0.986727 2000-01-02 1025 Edith -0.086358 -0.032920 2000-01-03 982 Edith 0.473177 0.298654 2000-01-04 1009 Sarah 0.534344 -0.750377 2000-01-05 963 Zelda -0.271573 0.054424 ... ... ... ... ... 2000-12-27 980 Ingrid -0.132333 -0.422195 2000-12-28 972 Frank -0.376007 -0.298687 2000-12-29 1009 Ursula -0.865047 -0.503133 2000-12-30 1000 Hannah -0.063757 -0.507336 2000-12-31 972 Tim -0.869120 0.531685 """ index = pd.date_range(start=start, end=end, freq=freq, name="timestamp") n = len(index) state = np.random.RandomState(seed) columns = { "name": state.choice(_names, size=n), "id": state.poisson(1000, size=n), "x": state.rand(n) * 2 - 1, "y": state.rand(n) * 2 - 1, } df = DataFrame(columns, index=index, columns=sorted(columns)) if df.index[-1] == end: df = df.iloc[:-1] return df def index_subclass_makers_generator(): make_index_funcs = [ makeDateIndex, makePeriodIndex, makeTimedeltaIndex, makeRangeIndex, makeIntervalIndex, makeCategoricalIndex, makeMultiIndex, ] yield from make_index_funcs def all_timeseries_index_generator(k: int = 10) -> Iterable[Index]: """ Generator which can be iterated over to get instances of all the classes which represent time-series. Parameters ---------- k: length of each of the index instances """ make_index_funcs: list[Callable[..., Index]] = [ makeDateIndex, makePeriodIndex, makeTimedeltaIndex, ] for make_index_func in make_index_funcs: yield make_index_func(k=k) # make series def makeFloatSeries(name=None): index = makeStringIndex(_N) return Series(np.random.randn(_N), index=index, name=name) def makeStringSeries(name=None): index = makeStringIndex(_N) return Series(np.random.randn(_N), index=index, name=name) def makeObjectSeries(name=None): data = makeStringIndex(_N) data = Index(data, dtype=object) index = makeStringIndex(_N) return Series(data, index=index, name=name) def getSeriesData(): index = makeStringIndex(_N) return {c: Series(np.random.randn(_N), index=index) for c in getCols(_K)} def makeTimeSeries(nper=None, freq="B", name=None): if nper is None: nper = _N return Series( np.random.randn(nper), index=makeDateIndex(nper, freq=freq), name=name ) def makePeriodSeries(nper=None, name=None): if nper is None: nper = _N return Series(np.random.randn(nper), index=makePeriodIndex(nper), name=name) def getTimeSeriesData(nper=None, freq="B"): return {c: makeTimeSeries(nper, freq) for c in getCols(_K)} def getPeriodData(nper=None): return {c: makePeriodSeries(nper) for c in getCols(_K)} # make frame def makeTimeDataFrame(nper=None, freq="B"): data = getTimeSeriesData(nper, freq) return DataFrame(data) def makeDataFrame() -> DataFrame: data = getSeriesData() return DataFrame(data) def getMixedTypeDict(): index = Index(["a", "b", "c", "d", "e"]) data = { "A": [0.0, 1.0, 2.0, 3.0, 4.0], "B": [0.0, 1.0, 0.0, 1.0, 0.0], "C": ["foo1", "foo2", "foo3", "foo4", "foo5"], "D": bdate_range("1/1/2009", periods=5), } return index, data def makeMixedDataFrame(): return DataFrame(getMixedTypeDict()[1]) def makePeriodFrame(nper=None): data = getPeriodData(nper) return DataFrame(data) def makeCustomIndex( nentries, nlevels, prefix="#", names=False, ndupe_l=None, idx_type=None ): """ Create an index/multindex with given dimensions, levels, names, etc' nentries - number of entries in index nlevels - number of levels (> 1 produces multindex) prefix - a string prefix for labels names - (Optional), bool or list of strings. if True will use default names, if false will use no names, if a list is given, the name of each level in the index will be taken from the list. ndupe_l - (Optional), list of ints, the number of rows for which the label will repeated at the corresponding level, you can specify just the first few, the rest will use the default ndupe_l of 1. len(ndupe_l) <= nlevels. idx_type - "i"/"f"/"s"/"u"/"dt"/"p"/"td". If idx_type is not None, `idx_nlevels` must be 1. "i"/"f" creates an integer/float index, "s"/"u" creates a string/unicode index "dt" create a datetime index. "td" create a datetime index. if unspecified, string labels will be generated. """ if ndupe_l is None: ndupe_l = [1] * nlevels assert is_sequence(ndupe_l) and len(ndupe_l) <= nlevels assert names is None or names is False or names is True or len(names) is nlevels assert idx_type is None or ( idx_type in ("i", "f", "s", "u", "dt", "p", "td") and nlevels == 1 ) if names is True: # build default names names = [prefix + str(i) for i in range(nlevels)] if names is False: # pass None to index constructor for no name names = None # make singleton case uniform if isinstance(names, str) and nlevels == 1: names = [names] # specific 1D index type requested? idx_func_dict: dict[str, Callable[..., Index]] = { "i": makeIntIndex, "f": makeFloatIndex, "s": makeStringIndex, "u": makeUnicodeIndex, "dt": makeDateIndex, "td": makeTimedeltaIndex, "p": makePeriodIndex, } idx_func = idx_func_dict.get(idx_type) if idx_func: idx = idx_func(nentries) # but we need to fill in the name if names: idx.name = names[0] return idx elif idx_type is not None: raise ValueError( f"{repr(idx_type)} is not a legal value for `idx_type`, " "use 'i'/'f'/'s'/'u'/'dt'/'p'/'td'." ) if len(ndupe_l) < nlevels: ndupe_l.extend([1] * (nlevels - len(ndupe_l))) assert len(ndupe_l) == nlevels assert all(x > 0 for x in ndupe_l) list_of_lists = [] for i in range(nlevels): def keyfunc(x): import re numeric_tuple = re.sub(r"[^\d_]_?", "", x).split("_") return [int(num) for num in numeric_tuple] # build a list of lists to create the index from div_factor = nentries // ndupe_l[i] + 1 # Deprecated since version 3.9: collections.Counter now supports []. See PEP 585 # and Generic Alias Type. cnt: Counter[str] = collections.Counter() for j in range(div_factor): label = f"{prefix}_l{i}_g{j}" cnt[label] = ndupe_l[i] # cute Counter trick result = sorted(cnt.elements(), key=keyfunc)[:nentries] list_of_lists.append(result) tuples = list(zip(*list_of_lists)) # convert tuples to index if nentries == 1: # we have a single level of tuples, i.e. a regular Index index = Index(tuples[0], name=names[0]) elif nlevels == 1: name = None if names is None else names[0] index = Index((x[0] for x in tuples), name=name) else: index = MultiIndex.from_tuples(tuples, names=names) return index def makeCustomDataframe( nrows, ncols, c_idx_names=True, r_idx_names=True, c_idx_nlevels=1, r_idx_nlevels=1, data_gen_f=None, c_ndupe_l=None, r_ndupe_l=None, dtype=None, c_idx_type=None, r_idx_type=None, ): """ Create a DataFrame using supplied parameters. Parameters ---------- nrows, ncols - number of data rows/cols c_idx_names, idx_names - False/True/list of strings, yields No names , default names or uses the provided names for the levels of the corresponding index. You can provide a single string when c_idx_nlevels ==1. c_idx_nlevels - number of levels in columns index. > 1 will yield MultiIndex r_idx_nlevels - number of levels in rows index. > 1 will yield MultiIndex data_gen_f - a function f(row,col) which return the data value at that position, the default generator used yields values of the form "RxCy" based on position. c_ndupe_l, r_ndupe_l - list of integers, determines the number of duplicates for each label at a given level of the corresponding index. The default `None` value produces a multiplicity of 1 across all levels, i.e. a unique index. Will accept a partial list of length N < idx_nlevels, for just the first N levels. If ndupe doesn't divide nrows/ncol, the last label might have lower multiplicity. dtype - passed to the DataFrame constructor as is, in case you wish to have more control in conjunction with a custom `data_gen_f` r_idx_type, c_idx_type - "i"/"f"/"s"/"u"/"dt"/"td". If idx_type is not None, `idx_nlevels` must be 1. "i"/"f" creates an integer/float index, "s"/"u" creates a string/unicode index "dt" create a datetime index. "td" create a timedelta index. if unspecified, string labels will be generated. Examples -------- # 5 row, 3 columns, default names on both, single index on both axis >> makeCustomDataframe(5,3) # make the data a random int between 1 and 100 >> mkdf(5,3,data_gen_f=lambda r,c:randint(1,100)) # 2-level multiindex on rows with each label duplicated # twice on first level, default names on both axis, single # index on both axis >> a=makeCustomDataframe(5,3,r_idx_nlevels=2,r_ndupe_l=[2]) # DatetimeIndex on row, index with unicode labels on columns # no names on either axis >> a=makeCustomDataframe(5,3,c_idx_names=False,r_idx_names=False, r_idx_type="dt",c_idx_type="u") # 4-level multindex on rows with names provided, 2-level multindex # on columns with default labels and default names. >> a=makeCustomDataframe(5,3,r_idx_nlevels=4, r_idx_names=["FEE","FIH","FOH","FUM"], c_idx_nlevels=2) >> a=mkdf(5,3,r_idx_nlevels=2,c_idx_nlevels=4) """ assert c_idx_nlevels > 0 assert r_idx_nlevels > 0 assert r_idx_type is None or ( r_idx_type in ("i", "f", "s", "u", "dt", "p", "td") and r_idx_nlevels == 1 ) assert c_idx_type is None or ( c_idx_type in ("i", "f", "s", "u", "dt", "p", "td") and c_idx_nlevels == 1 ) columns = makeCustomIndex( ncols, nlevels=c_idx_nlevels, prefix="C", names=c_idx_names, ndupe_l=c_ndupe_l, idx_type=c_idx_type, ) index = makeCustomIndex( nrows, nlevels=r_idx_nlevels, prefix="R", names=r_idx_names, ndupe_l=r_ndupe_l, idx_type=r_idx_type, ) # by default, generate data based on location if data_gen_f is None: data_gen_f = lambda r, c: f"R{r}C{c}" data = [[data_gen_f(r, c) for c in range(ncols)] for r in range(nrows)] return DataFrame(data, index, columns, dtype=dtype) def _create_missing_idx(nrows, ncols, density, random_state=None): if random_state is None: random_state = np.random else: random_state = np.random.RandomState(random_state) # below is cribbed from scipy.sparse size = round((1 - density) * nrows * ncols) # generate a few more to ensure unique values min_rows = 5 fac = 1.02 extra_size = min(size + min_rows, fac * size) def _gen_unique_rand(rng, _extra_size): ind = rng.rand(int(_extra_size)) return np.unique(np.floor(ind * nrows * ncols))[:size] ind = _gen_unique_rand(random_state, extra_size) while ind.size < size: extra_size *= 1.05 ind = _gen_unique_rand(random_state, extra_size) j = np.floor(ind * 1.0 / nrows).astype(int) i = (ind - j * nrows).astype(int) return i.tolist(), j.tolist() def makeMissingDataframe(density=0.9, random_state=None): df = makeDataFrame() i, j = _create_missing_idx(*df.shape, density=density, random_state=random_state) df.values[i, j] = np.nan return df def test_parallel(num_threads=2, kwargs_list=None): """ Decorator to run the same function multiple times in parallel. Parameters ---------- num_threads : int, optional The number of times the function is run in parallel. kwargs_list : list of dicts, optional The list of kwargs to update original function kwargs on different threads. Notes ----- This decorator does not pass the return value of the decorated function. Original from scikit-image: https://github.com/scikit-image/scikit-image/pull/1519 """ assert num_threads > 0 has_kwargs_list = kwargs_list is not None if has_kwargs_list: assert len(kwargs_list) == num_threads import threading def wrapper(func): @wraps(func) def inner(*args, **kwargs): if has_kwargs_list: update_kwargs = lambda i: dict(kwargs, **kwargs_list[i]) else: update_kwargs = lambda i: kwargs threads = [] for i in range(num_threads): updated_kwargs = update_kwargs(i) thread = threading.Thread(target=func, args=args, kwargs=updated_kwargs) threads.append(thread) for thread in threads: thread.start() for thread in threads: thread.join() return inner return wrapper class SubclassedSeries(Series): _metadata = ["testattr", "name"] @property def _constructor(self): return SubclassedSeries @property def _constructor_expanddim(self): return SubclassedDataFrame class SubclassedDataFrame(DataFrame): _metadata = ["testattr"] @property def _constructor(self): return SubclassedDataFrame @property def _constructor_sliced(self): return SubclassedSeries class SubclassedCategorical(Categorical): @property def _constructor(self): return SubclassedCategorical def _make_skipna_wrapper(alternative, skipna_alternative=None): """ Create a function for calling on an array. Parameters ---------- alternative : function The function to be called on the array with no NaNs. Only used when 'skipna_alternative' is None. skipna_alternative : function The function to be called on the original array Returns ------- function """ if skipna_alternative: def skipna_wrapper(x): return skipna_alternative(x.values) else: def skipna_wrapper(x): nona = x.dropna() if len(nona) == 0: return np.nan return alternative(nona) return skipna_wrapper def convert_rows_list_to_csv_str(rows_list: list[str]): """ Convert list of CSV rows to single CSV-formatted string for current OS. This method is used for creating expected value of to_csv() method. Parameters ---------- rows_list : List[str] Each element represents the row of csv. Returns ------- str Expected output of to_csv() in current OS. """ sep = os.linesep return sep.join(rows_list) + sep def external_error_raised(expected_exception: type[Exception]) -> ContextManager: """ Helper function to mark pytest.raises that have an external error message. Parameters ---------- expected_exception : Exception Expected error to raise. Returns ------- Callable Regular `pytest.raises` function with `match` equal to `None`. """ import pytest return pytest.raises(expected_exception, match=None) # noqa: PDF010 cython_table = pd.core.common._cython_table.items() def get_cython_table_params(ndframe, func_names_and_expected): """ Combine frame, functions from com._cython_table keys and expected result. Parameters ---------- ndframe : DataFrame or Series func_names_and_expected : Sequence of two items The first item is a name of a NDFrame method ('sum', 'prod') etc. The second item is the expected return value. Returns ------- list List of three items (DataFrame, function, expected result) """ results = [] for func_name, expected in func_names_and_expected: results.append((ndframe, func_name, expected)) results += [ (ndframe, func, expected) for func, name in cython_table if name == func_name ] return results def get_op_from_name(op_name: str) -> Callable: """ The operator function for a given op name. Parameters ---------- op_name : str The op name, in form of "add" or "__add__". Returns ------- function A function performing the operation. """ short_opname = op_name.strip("_") try: op = getattr(operator, short_opname) except AttributeError: # Assume it is the reverse operator rop = getattr(operator, short_opname[1:]) op = lambda x, y: rop(y, x) return op # ----------------------------------------------------------------------------- # Indexing test helpers def getitem(x): return x def setitem(x): return x def loc(x): return x.loc def iloc(x): return x.iloc def at(x): return x.at def iat(x): return x.iat # ----------------------------------------------------------------------------- def shares_memory(left, right) -> bool: """ Pandas-compat for np.shares_memory. """ if isinstance(left, np.ndarray) and isinstance(right, np.ndarray): return np.shares_memory(left, right) elif isinstance(left, np.ndarray): # Call with reversed args to get to unpacking logic below. return shares_memory(right, left) if isinstance(left, RangeIndex): return False if isinstance(left, MultiIndex): return shares_memory(left._codes, right) if isinstance(left, (Index, Series)): return shares_memory(left._values, right) if isinstance(left, NDArrayBackedExtensionArray): return shares_memory(left._ndarray, right) if isinstance(left, pd.core.arrays.SparseArray): return shares_memory(left.sp_values, right) if isinstance(left, pd.core.arrays.IntervalArray): return shares_memory(left._left, right) or shares_memory(left._right, right) if isinstance(left, ExtensionArray) and left.dtype == "string[pyarrow]": # https://github.com/pandas-dev/pandas/pull/43930#discussion_r736862669 if isinstance(right, ExtensionArray) and right.dtype == "string[pyarrow]": # error: "ExtensionArray" has no attribute "_data" left_pa_data = left._data # type: ignore[attr-defined] # error: "ExtensionArray" has no attribute "_data" right_pa_data = right._data # type: ignore[attr-defined] left_buf1 = left_pa_data.chunk(0).buffers()[1] right_buf1 = right_pa_data.chunk(0).buffers()[1] return left_buf1 == right_buf1 if isinstance(left, BaseMaskedArray) and isinstance(right, BaseMaskedArray): # By convention, we'll say these share memory if they share *either* # the _data or the _mask return np.shares_memory(left._data, right._data) or np.shares_memory( left._mask, right._mask ) if isinstance(left, DataFrame) and len(left._mgr.arrays) == 1: arr = left._mgr.arrays[0] return shares_memory(arr, right) raise NotImplementedError(type(left), type(right))