_testing.py 82 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755
  1. import bz2
  2. from collections import Counter
  3. from contextlib import contextmanager
  4. from datetime import datetime
  5. from functools import wraps
  6. import gzip
  7. import os
  8. from shutil import rmtree
  9. import string
  10. import tempfile
  11. from typing import Any, List, Optional, Union, cast
  12. import warnings
  13. import zipfile
  14. import numpy as np
  15. from numpy.random import rand, randn
  16. from pandas._config.localization import ( # noqa:F401
  17. can_set_locale,
  18. get_locales,
  19. set_locale,
  20. )
  21. import pandas._libs.testing as _testing
  22. from pandas._typing import FilePathOrBuffer, FrameOrSeries
  23. from pandas.compat import _get_lzma_file, _import_lzma
  24. from pandas.core.dtypes.common import (
  25. is_bool,
  26. is_categorical_dtype,
  27. is_datetime64_dtype,
  28. is_datetime64tz_dtype,
  29. is_extension_array_dtype,
  30. is_interval_dtype,
  31. is_list_like,
  32. is_number,
  33. is_period_dtype,
  34. is_sequence,
  35. is_timedelta64_dtype,
  36. needs_i8_conversion,
  37. )
  38. from pandas.core.dtypes.missing import array_equivalent
  39. import pandas as pd
  40. from pandas import (
  41. Categorical,
  42. CategoricalIndex,
  43. DataFrame,
  44. DatetimeIndex,
  45. Index,
  46. IntervalIndex,
  47. MultiIndex,
  48. RangeIndex,
  49. Series,
  50. bdate_range,
  51. )
  52. from pandas.core.algorithms import take_1d
  53. from pandas.core.arrays import (
  54. DatetimeArray,
  55. ExtensionArray,
  56. IntervalArray,
  57. PeriodArray,
  58. TimedeltaArray,
  59. period_array,
  60. )
  61. from pandas.io.common import urlopen
  62. from pandas.io.formats.printing import pprint_thing
  63. lzma = _import_lzma()
  64. N = 30
  65. K = 4
  66. _RAISE_NETWORK_ERROR_DEFAULT = False
  67. # set testing_mode
  68. _testing_mode_warnings = (DeprecationWarning, ResourceWarning)
  69. def set_testing_mode():
  70. # set the testing mode filters
  71. testing_mode = os.environ.get("PANDAS_TESTING_MODE", "None")
  72. if "deprecate" in testing_mode:
  73. warnings.simplefilter("always", _testing_mode_warnings)
  74. def reset_testing_mode():
  75. # reset the testing mode filters
  76. testing_mode = os.environ.get("PANDAS_TESTING_MODE", "None")
  77. if "deprecate" in testing_mode:
  78. warnings.simplefilter("ignore", _testing_mode_warnings)
  79. set_testing_mode()
  80. def reset_display_options():
  81. """
  82. Reset the display options for printing and representing objects.
  83. """
  84. pd.reset_option("^display.", silent=True)
  85. def round_trip_pickle(
  86. obj: Any, path: Optional[FilePathOrBuffer] = None
  87. ) -> FrameOrSeries:
  88. """
  89. Pickle an object and then read it again.
  90. Parameters
  91. ----------
  92. obj : any object
  93. The object to pickle and then re-read.
  94. path : str, path object or file-like object, default None
  95. The path where the pickled object is written and then read.
  96. Returns
  97. -------
  98. pandas object
  99. The original object that was pickled and then re-read.
  100. """
  101. _path = path
  102. if _path is None:
  103. _path = f"__{rands(10)}__.pickle"
  104. with ensure_clean(_path) as temp_path:
  105. pd.to_pickle(obj, temp_path)
  106. return pd.read_pickle(temp_path)
  107. def round_trip_pathlib(writer, reader, path: Optional[str] = None):
  108. """
  109. Write an object to file specified by a pathlib.Path and read it back
  110. Parameters
  111. ----------
  112. writer : callable bound to pandas object
  113. IO writing function (e.g. DataFrame.to_csv )
  114. reader : callable
  115. IO reading function (e.g. pd.read_csv )
  116. path : str, default None
  117. The path where the object is written and then read.
  118. Returns
  119. -------
  120. pandas object
  121. The original object that was serialized and then re-read.
  122. """
  123. import pytest
  124. Path = pytest.importorskip("pathlib").Path
  125. if path is None:
  126. path = "___pathlib___"
  127. with ensure_clean(path) as path:
  128. writer(Path(path))
  129. obj = reader(Path(path))
  130. return obj
  131. def round_trip_localpath(writer, reader, path: Optional[str] = None):
  132. """
  133. Write an object to file specified by a py.path LocalPath and read it back.
  134. Parameters
  135. ----------
  136. writer : callable bound to pandas object
  137. IO writing function (e.g. DataFrame.to_csv )
  138. reader : callable
  139. IO reading function (e.g. pd.read_csv )
  140. path : str, default None
  141. The path where the object is written and then read.
  142. Returns
  143. -------
  144. pandas object
  145. The original object that was serialized and then re-read.
  146. """
  147. import pytest
  148. LocalPath = pytest.importorskip("py.path").local
  149. if path is None:
  150. path = "___localpath___"
  151. with ensure_clean(path) as path:
  152. writer(LocalPath(path))
  153. obj = reader(LocalPath(path))
  154. return obj
  155. @contextmanager
  156. def decompress_file(path, compression):
  157. """
  158. Open a compressed file and return a file object.
  159. Parameters
  160. ----------
  161. path : str
  162. The path where the file is read from.
  163. compression : {'gzip', 'bz2', 'zip', 'xz', None}
  164. Name of the decompression to use
  165. Returns
  166. -------
  167. file object
  168. """
  169. if compression is None:
  170. f = open(path, "rb")
  171. elif compression == "gzip":
  172. f = gzip.open(path, "rb")
  173. elif compression == "bz2":
  174. f = bz2.BZ2File(path, "rb")
  175. elif compression == "xz":
  176. f = _get_lzma_file(lzma)(path, "rb")
  177. elif compression == "zip":
  178. zip_file = zipfile.ZipFile(path)
  179. zip_names = zip_file.namelist()
  180. if len(zip_names) == 1:
  181. f = zip_file.open(zip_names.pop())
  182. else:
  183. raise ValueError(f"ZIP file {path} error. Only one file per ZIP.")
  184. else:
  185. raise ValueError(f"Unrecognized compression type: {compression}")
  186. try:
  187. yield f
  188. finally:
  189. f.close()
  190. if compression == "zip":
  191. zip_file.close()
  192. def write_to_compressed(compression, path, data, dest="test"):
  193. """
  194. Write data to a compressed file.
  195. Parameters
  196. ----------
  197. compression : {'gzip', 'bz2', 'zip', 'xz'}
  198. The compression type to use.
  199. path : str
  200. The file path to write the data.
  201. data : str
  202. The data to write.
  203. dest : str, default "test"
  204. The destination file (for ZIP only)
  205. Raises
  206. ------
  207. ValueError : An invalid compression value was passed in.
  208. """
  209. if compression == "zip":
  210. import zipfile
  211. compress_method = zipfile.ZipFile
  212. elif compression == "gzip":
  213. import gzip
  214. compress_method = gzip.GzipFile
  215. elif compression == "bz2":
  216. import bz2
  217. compress_method = bz2.BZ2File
  218. elif compression == "xz":
  219. compress_method = _get_lzma_file(lzma)
  220. else:
  221. raise ValueError(f"Unrecognized compression type: {compression}")
  222. if compression == "zip":
  223. mode = "w"
  224. args = (dest, data)
  225. method = "writestr"
  226. else:
  227. mode = "wb"
  228. args = (data,)
  229. method = "write"
  230. with compress_method(path, mode=mode) as f:
  231. getattr(f, method)(*args)
  232. def assert_almost_equal(
  233. left,
  234. right,
  235. check_dtype: Union[bool, str] = "equiv",
  236. check_less_precise: Union[bool, int] = False,
  237. **kwargs,
  238. ):
  239. """
  240. Check that the left and right objects are approximately equal.
  241. By approximately equal, we refer to objects that are numbers or that
  242. contain numbers which may be equivalent to specific levels of precision.
  243. Parameters
  244. ----------
  245. left : object
  246. right : object
  247. check_dtype : bool or {'equiv'}, default 'equiv'
  248. Check dtype if both a and b are the same type. If 'equiv' is passed in,
  249. then `RangeIndex` and `Int64Index` are also considered equivalent
  250. when doing type checking.
  251. check_less_precise : bool or int, default False
  252. Specify comparison precision. 5 digits (False) or 3 digits (True)
  253. after decimal points are compared. If int, then specify the number
  254. of digits to compare.
  255. When comparing two numbers, if the first number has magnitude less
  256. than 1e-5, we compare the two numbers directly and check whether
  257. they are equivalent within the specified precision. Otherwise, we
  258. compare the **ratio** of the second number to the first number and
  259. check whether it is equivalent to 1 within the specified precision.
  260. """
  261. if isinstance(left, pd.Index):
  262. assert_index_equal(
  263. left,
  264. right,
  265. check_exact=False,
  266. exact=check_dtype,
  267. check_less_precise=check_less_precise,
  268. **kwargs,
  269. )
  270. elif isinstance(left, pd.Series):
  271. assert_series_equal(
  272. left,
  273. right,
  274. check_exact=False,
  275. check_dtype=check_dtype,
  276. check_less_precise=check_less_precise,
  277. **kwargs,
  278. )
  279. elif isinstance(left, pd.DataFrame):
  280. assert_frame_equal(
  281. left,
  282. right,
  283. check_exact=False,
  284. check_dtype=check_dtype,
  285. check_less_precise=check_less_precise,
  286. **kwargs,
  287. )
  288. else:
  289. # Other sequences.
  290. if check_dtype:
  291. if is_number(left) and is_number(right):
  292. # Do not compare numeric classes, like np.float64 and float.
  293. pass
  294. elif is_bool(left) and is_bool(right):
  295. # Do not compare bool classes, like np.bool_ and bool.
  296. pass
  297. else:
  298. if isinstance(left, np.ndarray) or isinstance(right, np.ndarray):
  299. obj = "numpy array"
  300. else:
  301. obj = "Input"
  302. assert_class_equal(left, right, obj=obj)
  303. _testing.assert_almost_equal(
  304. left,
  305. right,
  306. check_dtype=check_dtype,
  307. check_less_precise=check_less_precise,
  308. **kwargs,
  309. )
  310. def _check_isinstance(left, right, cls):
  311. """
  312. Helper method for our assert_* methods that ensures that
  313. the two objects being compared have the right type before
  314. proceeding with the comparison.
  315. Parameters
  316. ----------
  317. left : The first object being compared.
  318. right : The second object being compared.
  319. cls : The class type to check against.
  320. Raises
  321. ------
  322. AssertionError : Either `left` or `right` is not an instance of `cls`.
  323. """
  324. cls_name = cls.__name__
  325. if not isinstance(left, cls):
  326. raise AssertionError(
  327. f"{cls_name} Expected type {cls}, found {type(left)} instead"
  328. )
  329. if not isinstance(right, cls):
  330. raise AssertionError(
  331. f"{cls_name} Expected type {cls}, found {type(right)} instead"
  332. )
  333. def assert_dict_equal(left, right, compare_keys: bool = True):
  334. _check_isinstance(left, right, dict)
  335. _testing.assert_dict_equal(left, right, compare_keys=compare_keys)
  336. def randbool(size=(), p: float = 0.5):
  337. return rand(*size) <= p
  338. RANDS_CHARS = np.array(list(string.ascii_letters + string.digits), dtype=(np.str_, 1))
  339. RANDU_CHARS = np.array(
  340. list("".join(map(chr, range(1488, 1488 + 26))) + string.digits),
  341. dtype=(np.unicode_, 1),
  342. )
  343. def rands_array(nchars, size, dtype="O"):
  344. """
  345. Generate an array of byte strings.
  346. """
  347. retval = (
  348. np.random.choice(RANDS_CHARS, size=nchars * np.prod(size))
  349. .view((np.str_, nchars))
  350. .reshape(size)
  351. )
  352. if dtype is None:
  353. return retval
  354. else:
  355. return retval.astype(dtype)
  356. def randu_array(nchars, size, dtype="O"):
  357. """
  358. Generate an array of unicode strings.
  359. """
  360. retval = (
  361. np.random.choice(RANDU_CHARS, size=nchars * np.prod(size))
  362. .view((np.unicode_, nchars))
  363. .reshape(size)
  364. )
  365. if dtype is None:
  366. return retval
  367. else:
  368. return retval.astype(dtype)
  369. def rands(nchars):
  370. """
  371. Generate one random byte string.
  372. See `rands_array` if you want to create an array of random strings.
  373. """
  374. return "".join(np.random.choice(RANDS_CHARS, nchars))
  375. def randu(nchars):
  376. """
  377. Generate one random unicode string.
  378. See `randu_array` if you want to create an array of random unicode strings.
  379. """
  380. return "".join(np.random.choice(RANDU_CHARS, nchars))
  381. def close(fignum=None):
  382. from matplotlib.pyplot import get_fignums, close as _close
  383. if fignum is None:
  384. for fignum in get_fignums():
  385. _close(fignum)
  386. else:
  387. _close(fignum)
  388. # -----------------------------------------------------------------------------
  389. # contextmanager to ensure the file cleanup
  390. @contextmanager
  391. def ensure_clean(filename=None, return_filelike=False):
  392. """
  393. Gets a temporary path and agrees to remove on close.
  394. Parameters
  395. ----------
  396. filename : str (optional)
  397. if None, creates a temporary file which is then removed when out of
  398. scope. if passed, creates temporary file with filename as ending.
  399. return_filelike : bool (default False)
  400. if True, returns a file-like which is *always* cleaned. Necessary for
  401. savefig and other functions which want to append extensions.
  402. """
  403. filename = filename or ""
  404. fd = None
  405. if return_filelike:
  406. f = tempfile.TemporaryFile(suffix=filename)
  407. try:
  408. yield f
  409. finally:
  410. f.close()
  411. else:
  412. # don't generate tempfile if using a path with directory specified
  413. if len(os.path.dirname(filename)):
  414. raise ValueError("Can't pass a qualified name to ensure_clean()")
  415. try:
  416. fd, filename = tempfile.mkstemp(suffix=filename)
  417. except UnicodeEncodeError:
  418. import pytest
  419. pytest.skip("no unicode file names on this system")
  420. try:
  421. yield filename
  422. finally:
  423. try:
  424. os.close(fd)
  425. except OSError:
  426. print(f"Couldn't close file descriptor: {fd} (file: {filename})")
  427. try:
  428. if os.path.exists(filename):
  429. os.remove(filename)
  430. except OSError as e:
  431. print(f"Exception on removing file: {e}")
  432. @contextmanager
  433. def ensure_clean_dir():
  434. """
  435. Get a temporary directory path and agrees to remove on close.
  436. Yields
  437. ------
  438. Temporary directory path
  439. """
  440. directory_name = tempfile.mkdtemp(suffix="")
  441. try:
  442. yield directory_name
  443. finally:
  444. try:
  445. rmtree(directory_name)
  446. except OSError:
  447. pass
  448. @contextmanager
  449. def ensure_safe_environment_variables():
  450. """
  451. Get a context manager to safely set environment variables
  452. All changes will be undone on close, hence environment variables set
  453. within this contextmanager will neither persist nor change global state.
  454. """
  455. saved_environ = dict(os.environ)
  456. try:
  457. yield
  458. finally:
  459. os.environ.clear()
  460. os.environ.update(saved_environ)
  461. # -----------------------------------------------------------------------------
  462. # Comparators
  463. def equalContents(arr1, arr2) -> bool:
  464. """
  465. Checks if the set of unique elements of arr1 and arr2 are equivalent.
  466. """
  467. return frozenset(arr1) == frozenset(arr2)
  468. def assert_index_equal(
  469. left: Index,
  470. right: Index,
  471. exact: Union[bool, str] = "equiv",
  472. check_names: bool = True,
  473. check_less_precise: Union[bool, int] = False,
  474. check_exact: bool = True,
  475. check_categorical: bool = True,
  476. obj: str = "Index",
  477. ) -> None:
  478. """
  479. Check that left and right Index are equal.
  480. Parameters
  481. ----------
  482. left : Index
  483. right : Index
  484. exact : bool or {'equiv'}, default 'equiv'
  485. Whether to check the Index class, dtype and inferred_type
  486. are identical. If 'equiv', then RangeIndex can be substituted for
  487. Int64Index as well.
  488. check_names : bool, default True
  489. Whether to check the names attribute.
  490. check_less_precise : bool or int, default False
  491. Specify comparison precision. Only used when check_exact is False.
  492. 5 digits (False) or 3 digits (True) after decimal points are compared.
  493. If int, then specify the digits to compare.
  494. check_exact : bool, default True
  495. Whether to compare number exactly.
  496. check_categorical : bool, default True
  497. Whether to compare internal Categorical exactly.
  498. obj : str, default 'Index'
  499. Specify object name being compared, internally used to show appropriate
  500. assertion message.
  501. """
  502. __tracebackhide__ = True
  503. def _check_types(l, r, obj="Index"):
  504. if exact:
  505. assert_class_equal(l, r, exact=exact, obj=obj)
  506. # Skip exact dtype checking when `check_categorical` is False
  507. if check_categorical:
  508. assert_attr_equal("dtype", l, r, obj=obj)
  509. # allow string-like to have different inferred_types
  510. if l.inferred_type in ("string", "unicode"):
  511. assert r.inferred_type in ("string", "unicode")
  512. else:
  513. assert_attr_equal("inferred_type", l, r, obj=obj)
  514. def _get_ilevel_values(index, level):
  515. # accept level number only
  516. unique = index.levels[level]
  517. level_codes = index.codes[level]
  518. filled = take_1d(unique._values, level_codes, fill_value=unique._na_value)
  519. values = unique._shallow_copy(filled, name=index.names[level])
  520. return values
  521. # instance validation
  522. _check_isinstance(left, right, Index)
  523. # class / dtype comparison
  524. _check_types(left, right, obj=obj)
  525. # level comparison
  526. if left.nlevels != right.nlevels:
  527. msg1 = f"{obj} levels are different"
  528. msg2 = f"{left.nlevels}, {left}"
  529. msg3 = f"{right.nlevels}, {right}"
  530. raise_assert_detail(obj, msg1, msg2, msg3)
  531. # length comparison
  532. if len(left) != len(right):
  533. msg1 = f"{obj} length are different"
  534. msg2 = f"{len(left)}, {left}"
  535. msg3 = f"{len(right)}, {right}"
  536. raise_assert_detail(obj, msg1, msg2, msg3)
  537. # MultiIndex special comparison for little-friendly error messages
  538. if left.nlevels > 1:
  539. left = cast(MultiIndex, left)
  540. right = cast(MultiIndex, right)
  541. for level in range(left.nlevels):
  542. # cannot use get_level_values here because it can change dtype
  543. llevel = _get_ilevel_values(left, level)
  544. rlevel = _get_ilevel_values(right, level)
  545. lobj = f"MultiIndex level [{level}]"
  546. assert_index_equal(
  547. llevel,
  548. rlevel,
  549. exact=exact,
  550. check_names=check_names,
  551. check_less_precise=check_less_precise,
  552. check_exact=check_exact,
  553. obj=lobj,
  554. )
  555. # get_level_values may change dtype
  556. _check_types(left.levels[level], right.levels[level], obj=obj)
  557. # skip exact index checking when `check_categorical` is False
  558. if check_exact and check_categorical:
  559. if not left.equals(right):
  560. diff = np.sum((left.values != right.values).astype(int)) * 100.0 / len(left)
  561. msg = f"{obj} values are different ({np.round(diff, 5)} %)"
  562. raise_assert_detail(obj, msg, left, right)
  563. else:
  564. _testing.assert_almost_equal(
  565. left.values,
  566. right.values,
  567. check_less_precise=check_less_precise,
  568. check_dtype=exact,
  569. obj=obj,
  570. lobj=left,
  571. robj=right,
  572. )
  573. # metadata comparison
  574. if check_names:
  575. assert_attr_equal("names", left, right, obj=obj)
  576. if isinstance(left, pd.PeriodIndex) or isinstance(right, pd.PeriodIndex):
  577. assert_attr_equal("freq", left, right, obj=obj)
  578. if isinstance(left, pd.IntervalIndex) or isinstance(right, pd.IntervalIndex):
  579. assert_interval_array_equal(left.values, right.values)
  580. if check_categorical:
  581. if is_categorical_dtype(left) or is_categorical_dtype(right):
  582. assert_categorical_equal(left.values, right.values, obj=f"{obj} category")
  583. def assert_class_equal(left, right, exact: Union[bool, str] = True, obj="Input"):
  584. """
  585. Checks classes are equal.
  586. """
  587. __tracebackhide__ = True
  588. def repr_class(x):
  589. if isinstance(x, Index):
  590. # return Index as it is to include values in the error message
  591. return x
  592. try:
  593. return type(x).__name__
  594. except AttributeError:
  595. return repr(type(x))
  596. if exact == "equiv":
  597. if type(left) != type(right):
  598. # allow equivalence of Int64Index/RangeIndex
  599. types = {type(left).__name__, type(right).__name__}
  600. if len(types - {"Int64Index", "RangeIndex"}):
  601. msg = f"{obj} classes are not equivalent"
  602. raise_assert_detail(obj, msg, repr_class(left), repr_class(right))
  603. elif exact:
  604. if type(left) != type(right):
  605. msg = f"{obj} classes are different"
  606. raise_assert_detail(obj, msg, repr_class(left), repr_class(right))
  607. def assert_attr_equal(attr, left, right, obj="Attributes"):
  608. """checks attributes are equal. Both objects must have attribute.
  609. Parameters
  610. ----------
  611. attr : str
  612. Attribute name being compared.
  613. left : object
  614. right : object
  615. obj : str, default 'Attributes'
  616. Specify object name being compared, internally used to show appropriate
  617. assertion message
  618. """
  619. __tracebackhide__ = True
  620. left_attr = getattr(left, attr)
  621. right_attr = getattr(right, attr)
  622. if left_attr is right_attr:
  623. return True
  624. elif (
  625. is_number(left_attr)
  626. and np.isnan(left_attr)
  627. and is_number(right_attr)
  628. and np.isnan(right_attr)
  629. ):
  630. # np.nan
  631. return True
  632. try:
  633. result = left_attr == right_attr
  634. except TypeError:
  635. # datetimetz on rhs may raise TypeError
  636. result = False
  637. if not isinstance(result, bool):
  638. result = result.all()
  639. if result:
  640. return True
  641. else:
  642. msg = f'Attribute "{attr}" are different'
  643. raise_assert_detail(obj, msg, left_attr, right_attr)
  644. def assert_is_valid_plot_return_object(objs):
  645. import matplotlib.pyplot as plt
  646. if isinstance(objs, (pd.Series, np.ndarray)):
  647. for el in objs.ravel():
  648. msg = (
  649. "one of 'objs' is not a matplotlib Axes instance, "
  650. f"type encountered {repr(type(el).__name__)}"
  651. )
  652. assert isinstance(el, (plt.Axes, dict)), msg
  653. else:
  654. msg = (
  655. "objs is neither an ndarray of Artist instances nor a single "
  656. "ArtistArtist instance, tuple, or dict, 'objs' is a "
  657. f"{repr(type(objs).__name__)}"
  658. )
  659. assert isinstance(objs, (plt.Artist, tuple, dict)), msg
  660. def isiterable(obj):
  661. return hasattr(obj, "__iter__")
  662. def assert_is_sorted(seq):
  663. """Assert that the sequence is sorted."""
  664. if isinstance(seq, (Index, Series)):
  665. seq = seq.values
  666. # sorting does not change precisions
  667. assert_numpy_array_equal(seq, np.sort(np.array(seq)))
  668. def assert_categorical_equal(
  669. left, right, check_dtype=True, check_category_order=True, obj="Categorical"
  670. ):
  671. """Test that Categoricals are equivalent.
  672. Parameters
  673. ----------
  674. left : Categorical
  675. right : Categorical
  676. check_dtype : bool, default True
  677. Check that integer dtype of the codes are the same
  678. check_category_order : bool, default True
  679. Whether the order of the categories should be compared, which
  680. implies identical integer codes. If False, only the resulting
  681. values are compared. The ordered attribute is
  682. checked regardless.
  683. obj : str, default 'Categorical'
  684. Specify object name being compared, internally used to show appropriate
  685. assertion message
  686. """
  687. _check_isinstance(left, right, Categorical)
  688. if check_category_order:
  689. assert_index_equal(left.categories, right.categories, obj=f"{obj}.categories")
  690. assert_numpy_array_equal(
  691. left.codes, right.codes, check_dtype=check_dtype, obj=f"{obj}.codes",
  692. )
  693. else:
  694. assert_index_equal(
  695. left.categories.sort_values(),
  696. right.categories.sort_values(),
  697. obj=f"{obj}.categories",
  698. )
  699. assert_index_equal(
  700. left.categories.take(left.codes),
  701. right.categories.take(right.codes),
  702. obj=f"{obj}.values",
  703. )
  704. assert_attr_equal("ordered", left, right, obj=obj)
  705. def assert_interval_array_equal(left, right, exact="equiv", obj="IntervalArray"):
  706. """Test that two IntervalArrays are equivalent.
  707. Parameters
  708. ----------
  709. left, right : IntervalArray
  710. The IntervalArrays to compare.
  711. exact : bool or {'equiv'}, default 'equiv'
  712. Whether to check the Index class, dtype and inferred_type
  713. are identical. If 'equiv', then RangeIndex can be substituted for
  714. Int64Index as well.
  715. obj : str, default 'IntervalArray'
  716. Specify object name being compared, internally used to show appropriate
  717. assertion message
  718. """
  719. _check_isinstance(left, right, IntervalArray)
  720. assert_index_equal(left.left, right.left, exact=exact, obj=f"{obj}.left")
  721. assert_index_equal(left.right, right.right, exact=exact, obj=f"{obj}.left")
  722. assert_attr_equal("closed", left, right, obj=obj)
  723. def assert_period_array_equal(left, right, obj="PeriodArray"):
  724. _check_isinstance(left, right, PeriodArray)
  725. assert_numpy_array_equal(left._data, right._data, obj=f"{obj}.values")
  726. assert_attr_equal("freq", left, right, obj=obj)
  727. def assert_datetime_array_equal(left, right, obj="DatetimeArray"):
  728. __tracebackhide__ = True
  729. _check_isinstance(left, right, DatetimeArray)
  730. assert_numpy_array_equal(left._data, right._data, obj=f"{obj}._data")
  731. assert_attr_equal("freq", left, right, obj=obj)
  732. assert_attr_equal("tz", left, right, obj=obj)
  733. def assert_timedelta_array_equal(left, right, obj="TimedeltaArray"):
  734. __tracebackhide__ = True
  735. _check_isinstance(left, right, TimedeltaArray)
  736. assert_numpy_array_equal(left._data, right._data, obj=f"{obj}._data")
  737. assert_attr_equal("freq", left, right, obj=obj)
  738. def raise_assert_detail(obj, message, left, right, diff=None):
  739. __tracebackhide__ = True
  740. if isinstance(left, np.ndarray):
  741. left = pprint_thing(left)
  742. elif is_categorical_dtype(left):
  743. left = repr(left)
  744. if isinstance(right, np.ndarray):
  745. right = pprint_thing(right)
  746. elif is_categorical_dtype(right):
  747. right = repr(right)
  748. msg = f"""{obj} are different
  749. {message}
  750. [left]: {left}
  751. [right]: {right}"""
  752. if diff is not None:
  753. msg += f"\n[diff]: {diff}"
  754. raise AssertionError(msg)
  755. def assert_numpy_array_equal(
  756. left,
  757. right,
  758. strict_nan=False,
  759. check_dtype=True,
  760. err_msg=None,
  761. check_same=None,
  762. obj="numpy array",
  763. ):
  764. """
  765. Check that 'np.ndarray' is equivalent.
  766. Parameters
  767. ----------
  768. left, right : numpy.ndarray or iterable
  769. The two arrays to be compared.
  770. strict_nan : bool, default False
  771. If True, consider NaN and None to be different.
  772. check_dtype : bool, default True
  773. Check dtype if both a and b are np.ndarray.
  774. err_msg : str, default None
  775. If provided, used as assertion message.
  776. check_same : None|'copy'|'same', default None
  777. Ensure left and right refer/do not refer to the same memory area.
  778. obj : str, default 'numpy array'
  779. Specify object name being compared, internally used to show appropriate
  780. assertion message.
  781. """
  782. __tracebackhide__ = True
  783. # instance validation
  784. # Show a detailed error message when classes are different
  785. assert_class_equal(left, right, obj=obj)
  786. # both classes must be an np.ndarray
  787. _check_isinstance(left, right, np.ndarray)
  788. def _get_base(obj):
  789. return obj.base if getattr(obj, "base", None) is not None else obj
  790. left_base = _get_base(left)
  791. right_base = _get_base(right)
  792. if check_same == "same":
  793. if left_base is not right_base:
  794. raise AssertionError(f"{repr(left_base)} is not {repr(right_base)}")
  795. elif check_same == "copy":
  796. if left_base is right_base:
  797. raise AssertionError(f"{repr(left_base)} is {repr(right_base)}")
  798. def _raise(left, right, err_msg):
  799. if err_msg is None:
  800. if left.shape != right.shape:
  801. raise_assert_detail(
  802. obj, f"{obj} shapes are different", left.shape, right.shape,
  803. )
  804. diff = 0
  805. for l, r in zip(left, right):
  806. # count up differences
  807. if not array_equivalent(l, r, strict_nan=strict_nan):
  808. diff += 1
  809. diff = diff * 100.0 / left.size
  810. msg = f"{obj} values are different ({np.round(diff, 5)} %)"
  811. raise_assert_detail(obj, msg, left, right)
  812. raise AssertionError(err_msg)
  813. # compare shape and values
  814. if not array_equivalent(left, right, strict_nan=strict_nan):
  815. _raise(left, right, err_msg)
  816. if check_dtype:
  817. if isinstance(left, np.ndarray) and isinstance(right, np.ndarray):
  818. assert_attr_equal("dtype", left, right, obj=obj)
  819. def assert_extension_array_equal(
  820. left, right, check_dtype=True, check_less_precise=False, check_exact=False
  821. ):
  822. """Check that left and right ExtensionArrays are equal.
  823. Parameters
  824. ----------
  825. left, right : ExtensionArray
  826. The two arrays to compare
  827. check_dtype : bool, default True
  828. Whether to check if the ExtensionArray dtypes are identical.
  829. check_less_precise : bool or int, default False
  830. Specify comparison precision. Only used when check_exact is False.
  831. 5 digits (False) or 3 digits (True) after decimal points are compared.
  832. If int, then specify the digits to compare.
  833. check_exact : bool, default False
  834. Whether to compare number exactly.
  835. Notes
  836. -----
  837. Missing values are checked separately from valid values.
  838. A mask of missing values is computed for each and checked to match.
  839. The remaining all-valid values are cast to object dtype and checked.
  840. """
  841. assert isinstance(left, ExtensionArray), "left is not an ExtensionArray"
  842. assert isinstance(right, ExtensionArray), "right is not an ExtensionArray"
  843. if check_dtype:
  844. assert_attr_equal("dtype", left, right, obj="ExtensionArray")
  845. if hasattr(left, "asi8") and type(right) == type(left):
  846. # Avoid slow object-dtype comparisons
  847. assert_numpy_array_equal(left.asi8, right.asi8)
  848. return
  849. left_na = np.asarray(left.isna())
  850. right_na = np.asarray(right.isna())
  851. assert_numpy_array_equal(left_na, right_na, obj="ExtensionArray NA mask")
  852. left_valid = np.asarray(left[~left_na].astype(object))
  853. right_valid = np.asarray(right[~right_na].astype(object))
  854. if check_exact:
  855. assert_numpy_array_equal(left_valid, right_valid, obj="ExtensionArray")
  856. else:
  857. _testing.assert_almost_equal(
  858. left_valid,
  859. right_valid,
  860. check_dtype=check_dtype,
  861. check_less_precise=check_less_precise,
  862. obj="ExtensionArray",
  863. )
  864. # This could be refactored to use the NDFrame.equals method
  865. def assert_series_equal(
  866. left,
  867. right,
  868. check_dtype=True,
  869. check_index_type="equiv",
  870. check_series_type=True,
  871. check_less_precise=False,
  872. check_names=True,
  873. check_exact=False,
  874. check_datetimelike_compat=False,
  875. check_categorical=True,
  876. check_category_order=True,
  877. obj="Series",
  878. ):
  879. """
  880. Check that left and right Series are equal.
  881. Parameters
  882. ----------
  883. left : Series
  884. right : Series
  885. check_dtype : bool, default True
  886. Whether to check the Series dtype is identical.
  887. check_index_type : bool or {'equiv'}, default 'equiv'
  888. Whether to check the Index class, dtype and inferred_type
  889. are identical.
  890. check_series_type : bool, default True
  891. Whether to check the Series class is identical.
  892. check_less_precise : bool or int, default False
  893. Specify comparison precision. Only used when check_exact is False.
  894. 5 digits (False) or 3 digits (True) after decimal points are compared.
  895. If int, then specify the digits to compare.
  896. When comparing two numbers, if the first number has magnitude less
  897. than 1e-5, we compare the two numbers directly and check whether
  898. they are equivalent within the specified precision. Otherwise, we
  899. compare the **ratio** of the second number to the first number and
  900. check whether it is equivalent to 1 within the specified precision.
  901. check_names : bool, default True
  902. Whether to check the Series and Index names attribute.
  903. check_exact : bool, default False
  904. Whether to compare number exactly.
  905. check_datetimelike_compat : bool, default False
  906. Compare datetime-like which is comparable ignoring dtype.
  907. check_categorical : bool, default True
  908. Whether to compare internal Categorical exactly.
  909. check_category_order : bool, default True
  910. Whether to compare category order of internal Categoricals
  911. .. versionadded:: 1.0.2
  912. obj : str, default 'Series'
  913. Specify object name being compared, internally used to show appropriate
  914. assertion message.
  915. """
  916. __tracebackhide__ = True
  917. # instance validation
  918. _check_isinstance(left, right, Series)
  919. if check_series_type:
  920. # ToDo: There are some tests using rhs is sparse
  921. # lhs is dense. Should use assert_class_equal in future
  922. assert isinstance(left, type(right))
  923. # assert_class_equal(left, right, obj=obj)
  924. # length comparison
  925. if len(left) != len(right):
  926. msg1 = f"{len(left)}, {left.index}"
  927. msg2 = f"{len(right)}, {right.index}"
  928. raise_assert_detail(obj, "Series length are different", msg1, msg2)
  929. # index comparison
  930. assert_index_equal(
  931. left.index,
  932. right.index,
  933. exact=check_index_type,
  934. check_names=check_names,
  935. check_less_precise=check_less_precise,
  936. check_exact=check_exact,
  937. check_categorical=check_categorical,
  938. obj=f"{obj}.index",
  939. )
  940. if check_dtype:
  941. # We want to skip exact dtype checking when `check_categorical`
  942. # is False. We'll still raise if only one is a `Categorical`,
  943. # regardless of `check_categorical`
  944. if (
  945. is_categorical_dtype(left)
  946. and is_categorical_dtype(right)
  947. and not check_categorical
  948. ):
  949. pass
  950. else:
  951. assert_attr_equal("dtype", left, right, obj=f"Attributes of {obj}")
  952. if check_exact:
  953. assert_numpy_array_equal(
  954. left._internal_get_values(),
  955. right._internal_get_values(),
  956. check_dtype=check_dtype,
  957. obj=str(obj),
  958. )
  959. elif check_datetimelike_compat:
  960. # we want to check only if we have compat dtypes
  961. # e.g. integer and M|m are NOT compat, but we can simply check
  962. # the values in that case
  963. if needs_i8_conversion(left) or needs_i8_conversion(right):
  964. # datetimelike may have different objects (e.g. datetime.datetime
  965. # vs Timestamp) but will compare equal
  966. if not Index(left.values).equals(Index(right.values)):
  967. msg = (
  968. f"[datetimelike_compat=True] {left.values} "
  969. f"is not equal to {right.values}."
  970. )
  971. raise AssertionError(msg)
  972. else:
  973. assert_numpy_array_equal(
  974. left._internal_get_values(),
  975. right._internal_get_values(),
  976. check_dtype=check_dtype,
  977. )
  978. elif is_interval_dtype(left) or is_interval_dtype(right):
  979. assert_interval_array_equal(left.array, right.array)
  980. elif is_extension_array_dtype(left.dtype) and is_datetime64tz_dtype(left.dtype):
  981. # .values is an ndarray, but ._values is the ExtensionArray.
  982. # TODO: Use .array
  983. assert is_extension_array_dtype(right.dtype)
  984. assert_extension_array_equal(left._values, right._values)
  985. elif (
  986. is_extension_array_dtype(left)
  987. and not is_categorical_dtype(left)
  988. and is_extension_array_dtype(right)
  989. and not is_categorical_dtype(right)
  990. ):
  991. assert_extension_array_equal(left.array, right.array)
  992. else:
  993. _testing.assert_almost_equal(
  994. left._internal_get_values(),
  995. right._internal_get_values(),
  996. check_less_precise=check_less_precise,
  997. check_dtype=check_dtype,
  998. obj=str(obj),
  999. )
  1000. # metadata comparison
  1001. if check_names:
  1002. assert_attr_equal("name", left, right, obj=obj)
  1003. if check_categorical:
  1004. if is_categorical_dtype(left) or is_categorical_dtype(right):
  1005. assert_categorical_equal(
  1006. left.values,
  1007. right.values,
  1008. obj=f"{obj} category",
  1009. check_category_order=check_category_order,
  1010. )
  1011. # This could be refactored to use the NDFrame.equals method
  1012. def assert_frame_equal(
  1013. left,
  1014. right,
  1015. check_dtype=True,
  1016. check_index_type="equiv",
  1017. check_column_type="equiv",
  1018. check_frame_type=True,
  1019. check_less_precise=False,
  1020. check_names=True,
  1021. by_blocks=False,
  1022. check_exact=False,
  1023. check_datetimelike_compat=False,
  1024. check_categorical=True,
  1025. check_like=False,
  1026. obj="DataFrame",
  1027. ):
  1028. """
  1029. Check that left and right DataFrame are equal.
  1030. This function is intended to compare two DataFrames and output any
  1031. differences. Is is mostly intended for use in unit tests.
  1032. Additional parameters allow varying the strictness of the
  1033. equality checks performed.
  1034. Parameters
  1035. ----------
  1036. left : DataFrame
  1037. First DataFrame to compare.
  1038. right : DataFrame
  1039. Second DataFrame to compare.
  1040. check_dtype : bool, default True
  1041. Whether to check the DataFrame dtype is identical.
  1042. check_index_type : bool or {'equiv'}, default 'equiv'
  1043. Whether to check the Index class, dtype and inferred_type
  1044. are identical.
  1045. check_column_type : bool or {'equiv'}, default 'equiv'
  1046. Whether to check the columns class, dtype and inferred_type
  1047. are identical. Is passed as the ``exact`` argument of
  1048. :func:`assert_index_equal`.
  1049. check_frame_type : bool, default True
  1050. Whether to check the DataFrame class is identical.
  1051. check_less_precise : bool or int, default False
  1052. Specify comparison precision. Only used when check_exact is False.
  1053. 5 digits (False) or 3 digits (True) after decimal points are compared.
  1054. If int, then specify the digits to compare.
  1055. When comparing two numbers, if the first number has magnitude less
  1056. than 1e-5, we compare the two numbers directly and check whether
  1057. they are equivalent within the specified precision. Otherwise, we
  1058. compare the **ratio** of the second number to the first number and
  1059. check whether it is equivalent to 1 within the specified precision.
  1060. check_names : bool, default True
  1061. Whether to check that the `names` attribute for both the `index`
  1062. and `column` attributes of the DataFrame is identical.
  1063. by_blocks : bool, default False
  1064. Specify how to compare internal data. If False, compare by columns.
  1065. If True, compare by blocks.
  1066. check_exact : bool, default False
  1067. Whether to compare number exactly.
  1068. check_datetimelike_compat : bool, default False
  1069. Compare datetime-like which is comparable ignoring dtype.
  1070. check_categorical : bool, default True
  1071. Whether to compare internal Categorical exactly.
  1072. check_like : bool, default False
  1073. If True, ignore the order of index & columns.
  1074. Note: index labels must match their respective rows
  1075. (same as in columns) - same labels must be with the same data.
  1076. obj : str, default 'DataFrame'
  1077. Specify object name being compared, internally used to show appropriate
  1078. assertion message.
  1079. See Also
  1080. --------
  1081. assert_series_equal : Equivalent method for asserting Series equality.
  1082. DataFrame.equals : Check DataFrame equality.
  1083. Examples
  1084. --------
  1085. This example shows comparing two DataFrames that are equal
  1086. but with columns of differing dtypes.
  1087. >>> from pandas._testing import assert_frame_equal
  1088. >>> df1 = pd.DataFrame({'a': [1, 2], 'b': [3, 4]})
  1089. >>> df2 = pd.DataFrame({'a': [1, 2], 'b': [3.0, 4.0]})
  1090. df1 equals itself.
  1091. >>> assert_frame_equal(df1, df1)
  1092. df1 differs from df2 as column 'b' is of a different type.
  1093. >>> assert_frame_equal(df1, df2)
  1094. Traceback (most recent call last):
  1095. ...
  1096. AssertionError: Attributes of DataFrame.iloc[:, 1] (column name="b") are different
  1097. Attribute "dtype" are different
  1098. [left]: int64
  1099. [right]: float64
  1100. Ignore differing dtypes in columns with check_dtype.
  1101. >>> assert_frame_equal(df1, df2, check_dtype=False)
  1102. """
  1103. __tracebackhide__ = True
  1104. # instance validation
  1105. _check_isinstance(left, right, DataFrame)
  1106. if check_frame_type:
  1107. assert isinstance(left, type(right))
  1108. # assert_class_equal(left, right, obj=obj)
  1109. # shape comparison
  1110. if left.shape != right.shape:
  1111. raise_assert_detail(
  1112. obj, f"{obj} shape mismatch", f"{repr(left.shape)}", f"{repr(right.shape)}",
  1113. )
  1114. if check_like:
  1115. left, right = left.reindex_like(right), right
  1116. # index comparison
  1117. assert_index_equal(
  1118. left.index,
  1119. right.index,
  1120. exact=check_index_type,
  1121. check_names=check_names,
  1122. check_less_precise=check_less_precise,
  1123. check_exact=check_exact,
  1124. check_categorical=check_categorical,
  1125. obj=f"{obj}.index",
  1126. )
  1127. # column comparison
  1128. assert_index_equal(
  1129. left.columns,
  1130. right.columns,
  1131. exact=check_column_type,
  1132. check_names=check_names,
  1133. check_less_precise=check_less_precise,
  1134. check_exact=check_exact,
  1135. check_categorical=check_categorical,
  1136. obj=f"{obj}.columns",
  1137. )
  1138. # compare by blocks
  1139. if by_blocks:
  1140. rblocks = right._to_dict_of_blocks()
  1141. lblocks = left._to_dict_of_blocks()
  1142. for dtype in list(set(list(lblocks.keys()) + list(rblocks.keys()))):
  1143. assert dtype in lblocks
  1144. assert dtype in rblocks
  1145. assert_frame_equal(
  1146. lblocks[dtype], rblocks[dtype], check_dtype=check_dtype, obj=obj
  1147. )
  1148. # compare by columns
  1149. else:
  1150. for i, col in enumerate(left.columns):
  1151. assert col in right
  1152. lcol = left.iloc[:, i]
  1153. rcol = right.iloc[:, i]
  1154. assert_series_equal(
  1155. lcol,
  1156. rcol,
  1157. check_dtype=check_dtype,
  1158. check_index_type=check_index_type,
  1159. check_less_precise=check_less_precise,
  1160. check_exact=check_exact,
  1161. check_names=check_names,
  1162. check_datetimelike_compat=check_datetimelike_compat,
  1163. check_categorical=check_categorical,
  1164. obj=f'{obj}.iloc[:, {i}] (column name="{col}")',
  1165. )
  1166. def assert_equal(left, right, **kwargs):
  1167. """
  1168. Wrapper for tm.assert_*_equal to dispatch to the appropriate test function.
  1169. Parameters
  1170. ----------
  1171. left, right : Index, Series, DataFrame, ExtensionArray, or np.ndarray
  1172. The two items to be compared.
  1173. **kwargs
  1174. All keyword arguments are passed through to the underlying assert method.
  1175. """
  1176. __tracebackhide__ = True
  1177. if isinstance(left, pd.Index):
  1178. assert_index_equal(left, right, **kwargs)
  1179. elif isinstance(left, pd.Series):
  1180. assert_series_equal(left, right, **kwargs)
  1181. elif isinstance(left, pd.DataFrame):
  1182. assert_frame_equal(left, right, **kwargs)
  1183. elif isinstance(left, IntervalArray):
  1184. assert_interval_array_equal(left, right, **kwargs)
  1185. elif isinstance(left, PeriodArray):
  1186. assert_period_array_equal(left, right, **kwargs)
  1187. elif isinstance(left, DatetimeArray):
  1188. assert_datetime_array_equal(left, right, **kwargs)
  1189. elif isinstance(left, TimedeltaArray):
  1190. assert_timedelta_array_equal(left, right, **kwargs)
  1191. elif isinstance(left, ExtensionArray):
  1192. assert_extension_array_equal(left, right, **kwargs)
  1193. elif isinstance(left, np.ndarray):
  1194. assert_numpy_array_equal(left, right, **kwargs)
  1195. elif isinstance(left, str):
  1196. assert kwargs == {}
  1197. assert left == right
  1198. else:
  1199. raise NotImplementedError(type(left))
  1200. def box_expected(expected, box_cls, transpose=True):
  1201. """
  1202. Helper function to wrap the expected output of a test in a given box_class.
  1203. Parameters
  1204. ----------
  1205. expected : np.ndarray, Index, Series
  1206. box_cls : {Index, Series, DataFrame}
  1207. Returns
  1208. -------
  1209. subclass of box_cls
  1210. """
  1211. if box_cls is pd.Index:
  1212. expected = pd.Index(expected)
  1213. elif box_cls is pd.Series:
  1214. expected = pd.Series(expected)
  1215. elif box_cls is pd.DataFrame:
  1216. expected = pd.Series(expected).to_frame()
  1217. if transpose:
  1218. # for vector operations, we we need a DataFrame to be a single-row,
  1219. # not a single-column, in order to operate against non-DataFrame
  1220. # vectors of the same length.
  1221. expected = expected.T
  1222. elif box_cls is PeriodArray:
  1223. # the PeriodArray constructor is not as flexible as period_array
  1224. expected = period_array(expected)
  1225. elif box_cls is DatetimeArray:
  1226. expected = DatetimeArray(expected)
  1227. elif box_cls is TimedeltaArray:
  1228. expected = TimedeltaArray(expected)
  1229. elif box_cls is np.ndarray:
  1230. expected = np.array(expected)
  1231. elif box_cls is to_array:
  1232. expected = to_array(expected)
  1233. else:
  1234. raise NotImplementedError(box_cls)
  1235. return expected
  1236. def to_array(obj):
  1237. # temporary implementation until we get pd.array in place
  1238. if is_period_dtype(obj):
  1239. return period_array(obj)
  1240. elif is_datetime64_dtype(obj) or is_datetime64tz_dtype(obj):
  1241. return DatetimeArray._from_sequence(obj)
  1242. elif is_timedelta64_dtype(obj):
  1243. return TimedeltaArray._from_sequence(obj)
  1244. else:
  1245. return np.array(obj)
  1246. # -----------------------------------------------------------------------------
  1247. # Sparse
  1248. def assert_sp_array_equal(
  1249. left,
  1250. right,
  1251. check_dtype=True,
  1252. check_kind=True,
  1253. check_fill_value=True,
  1254. consolidate_block_indices=False,
  1255. ):
  1256. """Check that the left and right SparseArray are equal.
  1257. Parameters
  1258. ----------
  1259. left : SparseArray
  1260. right : SparseArray
  1261. check_dtype : bool, default True
  1262. Whether to check the data dtype is identical.
  1263. check_kind : bool, default True
  1264. Whether to just the kind of the sparse index for each column.
  1265. check_fill_value : bool, default True
  1266. Whether to check that left.fill_value matches right.fill_value
  1267. consolidate_block_indices : bool, default False
  1268. Whether to consolidate contiguous blocks for sparse arrays with
  1269. a BlockIndex. Some operations, e.g. concat, will end up with
  1270. block indices that could be consolidated. Setting this to true will
  1271. create a new BlockIndex for that array, with consolidated
  1272. block indices.
  1273. """
  1274. _check_isinstance(left, right, pd.arrays.SparseArray)
  1275. assert_numpy_array_equal(left.sp_values, right.sp_values, check_dtype=check_dtype)
  1276. # SparseIndex comparison
  1277. assert isinstance(left.sp_index, pd._libs.sparse.SparseIndex)
  1278. assert isinstance(right.sp_index, pd._libs.sparse.SparseIndex)
  1279. if not check_kind:
  1280. left_index = left.sp_index.to_block_index()
  1281. right_index = right.sp_index.to_block_index()
  1282. else:
  1283. left_index = left.sp_index
  1284. right_index = right.sp_index
  1285. if consolidate_block_indices and left.kind == "block":
  1286. # we'll probably remove this hack...
  1287. left_index = left_index.to_int_index().to_block_index()
  1288. right_index = right_index.to_int_index().to_block_index()
  1289. if not left_index.equals(right_index):
  1290. raise_assert_detail(
  1291. "SparseArray.index", "index are not equal", left_index, right_index
  1292. )
  1293. else:
  1294. # Just ensure a
  1295. pass
  1296. if check_fill_value:
  1297. assert_attr_equal("fill_value", left, right)
  1298. if check_dtype:
  1299. assert_attr_equal("dtype", left, right)
  1300. assert_numpy_array_equal(left.to_dense(), right.to_dense(), check_dtype=check_dtype)
  1301. # -----------------------------------------------------------------------------
  1302. # Others
  1303. def assert_contains_all(iterable, dic):
  1304. for k in iterable:
  1305. assert k in dic, f"Did not contain item: {repr(k)}"
  1306. def assert_copy(iter1, iter2, **eql_kwargs):
  1307. """
  1308. iter1, iter2: iterables that produce elements
  1309. comparable with assert_almost_equal
  1310. Checks that the elements are equal, but not
  1311. the same object. (Does not check that items
  1312. in sequences are also not the same object)
  1313. """
  1314. for elem1, elem2 in zip(iter1, iter2):
  1315. assert_almost_equal(elem1, elem2, **eql_kwargs)
  1316. msg = (
  1317. f"Expected object {repr(type(elem1))} and object {repr(type(elem2))} to be "
  1318. "different objects, but they were the same object."
  1319. )
  1320. assert elem1 is not elem2, msg
  1321. def getCols(k):
  1322. return string.ascii_uppercase[:k]
  1323. # make index
  1324. def makeStringIndex(k=10, name=None):
  1325. return Index(rands_array(nchars=10, size=k), name=name)
  1326. def makeUnicodeIndex(k=10, name=None):
  1327. return Index(randu_array(nchars=10, size=k), name=name)
  1328. def makeCategoricalIndex(k=10, n=3, name=None, **kwargs):
  1329. """ make a length k index or n categories """
  1330. x = rands_array(nchars=4, size=n)
  1331. return CategoricalIndex(
  1332. Categorical.from_codes(np.arange(k) % n, categories=x), name=name, **kwargs
  1333. )
  1334. def makeIntervalIndex(k=10, name=None, **kwargs):
  1335. """ make a length k IntervalIndex """
  1336. x = np.linspace(0, 100, num=(k + 1))
  1337. return IntervalIndex.from_breaks(x, name=name, **kwargs)
  1338. def makeBoolIndex(k=10, name=None):
  1339. if k == 1:
  1340. return Index([True], name=name)
  1341. elif k == 2:
  1342. return Index([False, True], name=name)
  1343. return Index([False, True] + [False] * (k - 2), name=name)
  1344. def makeIntIndex(k=10, name=None):
  1345. return Index(list(range(k)), name=name)
  1346. def makeUIntIndex(k=10, name=None):
  1347. return Index([2 ** 63 + i for i in range(k)], name=name)
  1348. def makeRangeIndex(k=10, name=None, **kwargs):
  1349. return RangeIndex(0, k, 1, name=name, **kwargs)
  1350. def makeFloatIndex(k=10, name=None):
  1351. values = sorted(np.random.random_sample(k)) - np.random.random_sample(1)
  1352. return Index(values * (10 ** np.random.randint(0, 9)), name=name)
  1353. def makeDateIndex(k=10, freq="B", name=None, **kwargs):
  1354. dt = datetime(2000, 1, 1)
  1355. dr = bdate_range(dt, periods=k, freq=freq, name=name)
  1356. return DatetimeIndex(dr, name=name, **kwargs)
  1357. def makeTimedeltaIndex(k=10, freq="D", name=None, **kwargs):
  1358. return pd.timedelta_range(start="1 day", periods=k, freq=freq, name=name, **kwargs)
  1359. def makePeriodIndex(k=10, name=None, **kwargs):
  1360. dt = datetime(2000, 1, 1)
  1361. dr = pd.period_range(start=dt, periods=k, freq="B", name=name, **kwargs)
  1362. return dr
  1363. def makeMultiIndex(k=10, names=None, **kwargs):
  1364. return MultiIndex.from_product((("foo", "bar"), (1, 2)), names=names, **kwargs)
  1365. _names = [
  1366. "Alice",
  1367. "Bob",
  1368. "Charlie",
  1369. "Dan",
  1370. "Edith",
  1371. "Frank",
  1372. "George",
  1373. "Hannah",
  1374. "Ingrid",
  1375. "Jerry",
  1376. "Kevin",
  1377. "Laura",
  1378. "Michael",
  1379. "Norbert",
  1380. "Oliver",
  1381. "Patricia",
  1382. "Quinn",
  1383. "Ray",
  1384. "Sarah",
  1385. "Tim",
  1386. "Ursula",
  1387. "Victor",
  1388. "Wendy",
  1389. "Xavier",
  1390. "Yvonne",
  1391. "Zelda",
  1392. ]
  1393. def _make_timeseries(start="2000-01-01", end="2000-12-31", freq="1D", seed=None):
  1394. """
  1395. Make a DataFrame with a DatetimeIndex
  1396. Parameters
  1397. ----------
  1398. start : str or Timestamp, default "2000-01-01"
  1399. The start of the index. Passed to date_range with `freq`.
  1400. end : str or Timestamp, default "2000-12-31"
  1401. The end of the index. Passed to date_range with `freq`.
  1402. freq : str or Freq
  1403. The frequency to use for the DatetimeIndex
  1404. seed : int, optional
  1405. The random state seed.
  1406. * name : object dtype with string names
  1407. * id : int dtype with
  1408. * x, y : float dtype
  1409. Examples
  1410. --------
  1411. >>> _make_timeseries()
  1412. id name x y
  1413. timestamp
  1414. 2000-01-01 982 Frank 0.031261 0.986727
  1415. 2000-01-02 1025 Edith -0.086358 -0.032920
  1416. 2000-01-03 982 Edith 0.473177 0.298654
  1417. 2000-01-04 1009 Sarah 0.534344 -0.750377
  1418. 2000-01-05 963 Zelda -0.271573 0.054424
  1419. ... ... ... ... ...
  1420. 2000-12-27 980 Ingrid -0.132333 -0.422195
  1421. 2000-12-28 972 Frank -0.376007 -0.298687
  1422. 2000-12-29 1009 Ursula -0.865047 -0.503133
  1423. 2000-12-30 1000 Hannah -0.063757 -0.507336
  1424. 2000-12-31 972 Tim -0.869120 0.531685
  1425. """
  1426. index = pd.date_range(start=start, end=end, freq=freq, name="timestamp")
  1427. n = len(index)
  1428. state = np.random.RandomState(seed)
  1429. columns = {
  1430. "name": state.choice(_names, size=n),
  1431. "id": state.poisson(1000, size=n),
  1432. "x": state.rand(n) * 2 - 1,
  1433. "y": state.rand(n) * 2 - 1,
  1434. }
  1435. df = pd.DataFrame(columns, index=index, columns=sorted(columns))
  1436. if df.index[-1] == end:
  1437. df = df.iloc[:-1]
  1438. return df
  1439. def all_index_generator(k=10):
  1440. """Generator which can be iterated over to get instances of all the various
  1441. index classes.
  1442. Parameters
  1443. ----------
  1444. k: length of each of the index instances
  1445. """
  1446. all_make_index_funcs = [
  1447. makeIntIndex,
  1448. makeFloatIndex,
  1449. makeStringIndex,
  1450. makeUnicodeIndex,
  1451. makeDateIndex,
  1452. makePeriodIndex,
  1453. makeTimedeltaIndex,
  1454. makeBoolIndex,
  1455. makeRangeIndex,
  1456. makeIntervalIndex,
  1457. makeCategoricalIndex,
  1458. ]
  1459. for make_index_func in all_make_index_funcs:
  1460. yield make_index_func(k=k)
  1461. def index_subclass_makers_generator():
  1462. make_index_funcs = [
  1463. makeDateIndex,
  1464. makePeriodIndex,
  1465. makeTimedeltaIndex,
  1466. makeRangeIndex,
  1467. makeIntervalIndex,
  1468. makeCategoricalIndex,
  1469. makeMultiIndex,
  1470. ]
  1471. for make_index_func in make_index_funcs:
  1472. yield make_index_func
  1473. def all_timeseries_index_generator(k=10):
  1474. """Generator which can be iterated over to get instances of all the classes
  1475. which represent time-series.
  1476. Parameters
  1477. ----------
  1478. k: length of each of the index instances
  1479. """
  1480. make_index_funcs = [makeDateIndex, makePeriodIndex, makeTimedeltaIndex]
  1481. for make_index_func in make_index_funcs:
  1482. yield make_index_func(k=k)
  1483. # make series
  1484. def makeFloatSeries(name=None):
  1485. index = makeStringIndex(N)
  1486. return Series(randn(N), index=index, name=name)
  1487. def makeStringSeries(name=None):
  1488. index = makeStringIndex(N)
  1489. return Series(randn(N), index=index, name=name)
  1490. def makeObjectSeries(name=None):
  1491. data = makeStringIndex(N)
  1492. data = Index(data, dtype=object)
  1493. index = makeStringIndex(N)
  1494. return Series(data, index=index, name=name)
  1495. def getSeriesData():
  1496. index = makeStringIndex(N)
  1497. return {c: Series(randn(N), index=index) for c in getCols(K)}
  1498. def makeTimeSeries(nper=None, freq="B", name=None):
  1499. if nper is None:
  1500. nper = N
  1501. return Series(randn(nper), index=makeDateIndex(nper, freq=freq), name=name)
  1502. def makePeriodSeries(nper=None, name=None):
  1503. if nper is None:
  1504. nper = N
  1505. return Series(randn(nper), index=makePeriodIndex(nper), name=name)
  1506. def getTimeSeriesData(nper=None, freq="B"):
  1507. return {c: makeTimeSeries(nper, freq) for c in getCols(K)}
  1508. def getPeriodData(nper=None):
  1509. return {c: makePeriodSeries(nper) for c in getCols(K)}
  1510. # make frame
  1511. def makeTimeDataFrame(nper=None, freq="B"):
  1512. data = getTimeSeriesData(nper, freq)
  1513. return DataFrame(data)
  1514. def makeDataFrame():
  1515. data = getSeriesData()
  1516. return DataFrame(data)
  1517. def getMixedTypeDict():
  1518. index = Index(["a", "b", "c", "d", "e"])
  1519. data = {
  1520. "A": [0.0, 1.0, 2.0, 3.0, 4.0],
  1521. "B": [0.0, 1.0, 0.0, 1.0, 0.0],
  1522. "C": ["foo1", "foo2", "foo3", "foo4", "foo5"],
  1523. "D": bdate_range("1/1/2009", periods=5),
  1524. }
  1525. return index, data
  1526. def makeMixedDataFrame():
  1527. return DataFrame(getMixedTypeDict()[1])
  1528. def makePeriodFrame(nper=None):
  1529. data = getPeriodData(nper)
  1530. return DataFrame(data)
  1531. def makeCustomIndex(
  1532. nentries, nlevels, prefix="#", names=False, ndupe_l=None, idx_type=None
  1533. ):
  1534. """Create an index/multindex with given dimensions, levels, names, etc'
  1535. nentries - number of entries in index
  1536. nlevels - number of levels (> 1 produces multindex)
  1537. prefix - a string prefix for labels
  1538. names - (Optional), bool or list of strings. if True will use default
  1539. names, if false will use no names, if a list is given, the name of
  1540. each level in the index will be taken from the list.
  1541. ndupe_l - (Optional), list of ints, the number of rows for which the
  1542. label will repeated at the corresponding level, you can specify just
  1543. the first few, the rest will use the default ndupe_l of 1.
  1544. len(ndupe_l) <= nlevels.
  1545. idx_type - "i"/"f"/"s"/"u"/"dt"/"p"/"td".
  1546. If idx_type is not None, `idx_nlevels` must be 1.
  1547. "i"/"f" creates an integer/float index,
  1548. "s"/"u" creates a string/unicode index
  1549. "dt" create a datetime index.
  1550. "td" create a datetime index.
  1551. if unspecified, string labels will be generated.
  1552. """
  1553. if ndupe_l is None:
  1554. ndupe_l = [1] * nlevels
  1555. assert is_sequence(ndupe_l) and len(ndupe_l) <= nlevels
  1556. assert names is None or names is False or names is True or len(names) is nlevels
  1557. assert idx_type is None or (
  1558. idx_type in ("i", "f", "s", "u", "dt", "p", "td") and nlevels == 1
  1559. )
  1560. if names is True:
  1561. # build default names
  1562. names = [prefix + str(i) for i in range(nlevels)]
  1563. if names is False:
  1564. # pass None to index constructor for no name
  1565. names = None
  1566. # make singleton case uniform
  1567. if isinstance(names, str) and nlevels == 1:
  1568. names = [names]
  1569. # specific 1D index type requested?
  1570. idx_func = dict(
  1571. i=makeIntIndex,
  1572. f=makeFloatIndex,
  1573. s=makeStringIndex,
  1574. u=makeUnicodeIndex,
  1575. dt=makeDateIndex,
  1576. td=makeTimedeltaIndex,
  1577. p=makePeriodIndex,
  1578. ).get(idx_type)
  1579. if idx_func:
  1580. idx = idx_func(nentries)
  1581. # but we need to fill in the name
  1582. if names:
  1583. idx.name = names[0]
  1584. return idx
  1585. elif idx_type is not None:
  1586. raise ValueError(
  1587. f"{repr(idx_type)} is not a legal value for `idx_type`, "
  1588. "use 'i'/'f'/'s'/'u'/'dt'/'p'/'td'."
  1589. )
  1590. if len(ndupe_l) < nlevels:
  1591. ndupe_l.extend([1] * (nlevels - len(ndupe_l)))
  1592. assert len(ndupe_l) == nlevels
  1593. assert all(x > 0 for x in ndupe_l)
  1594. tuples = []
  1595. for i in range(nlevels):
  1596. def keyfunc(x):
  1597. import re
  1598. numeric_tuple = re.sub(r"[^\d_]_?", "", x).split("_")
  1599. return [int(num) for num in numeric_tuple]
  1600. # build a list of lists to create the index from
  1601. div_factor = nentries // ndupe_l[i] + 1
  1602. cnt = Counter()
  1603. for j in range(div_factor):
  1604. label = f"{prefix}_l{i}_g{j}"
  1605. cnt[label] = ndupe_l[i]
  1606. # cute Counter trick
  1607. result = sorted(cnt.elements(), key=keyfunc)[:nentries]
  1608. tuples.append(result)
  1609. tuples = list(zip(*tuples))
  1610. # convert tuples to index
  1611. if nentries == 1:
  1612. # we have a single level of tuples, i.e. a regular Index
  1613. index = Index(tuples[0], name=names[0])
  1614. elif nlevels == 1:
  1615. name = None if names is None else names[0]
  1616. index = Index((x[0] for x in tuples), name=name)
  1617. else:
  1618. index = MultiIndex.from_tuples(tuples, names=names)
  1619. return index
  1620. def makeCustomDataframe(
  1621. nrows,
  1622. ncols,
  1623. c_idx_names=True,
  1624. r_idx_names=True,
  1625. c_idx_nlevels=1,
  1626. r_idx_nlevels=1,
  1627. data_gen_f=None,
  1628. c_ndupe_l=None,
  1629. r_ndupe_l=None,
  1630. dtype=None,
  1631. c_idx_type=None,
  1632. r_idx_type=None,
  1633. ):
  1634. """
  1635. nrows, ncols - number of data rows/cols
  1636. c_idx_names, idx_names - False/True/list of strings, yields No names ,
  1637. default names or uses the provided names for the levels of the
  1638. corresponding index. You can provide a single string when
  1639. c_idx_nlevels ==1.
  1640. c_idx_nlevels - number of levels in columns index. > 1 will yield MultiIndex
  1641. r_idx_nlevels - number of levels in rows index. > 1 will yield MultiIndex
  1642. data_gen_f - a function f(row,col) which return the data value
  1643. at that position, the default generator used yields values of the form
  1644. "RxCy" based on position.
  1645. c_ndupe_l, r_ndupe_l - list of integers, determines the number
  1646. of duplicates for each label at a given level of the corresponding
  1647. index. The default `None` value produces a multiplicity of 1 across
  1648. all levels, i.e. a unique index. Will accept a partial list of length
  1649. N < idx_nlevels, for just the first N levels. If ndupe doesn't divide
  1650. nrows/ncol, the last label might have lower multiplicity.
  1651. dtype - passed to the DataFrame constructor as is, in case you wish to
  1652. have more control in conjunction with a custom `data_gen_f`
  1653. r_idx_type, c_idx_type - "i"/"f"/"s"/"u"/"dt"/"td".
  1654. If idx_type is not None, `idx_nlevels` must be 1.
  1655. "i"/"f" creates an integer/float index,
  1656. "s"/"u" creates a string/unicode index
  1657. "dt" create a datetime index.
  1658. "td" create a timedelta index.
  1659. if unspecified, string labels will be generated.
  1660. Examples:
  1661. # 5 row, 3 columns, default names on both, single index on both axis
  1662. >> makeCustomDataframe(5,3)
  1663. # make the data a random int between 1 and 100
  1664. >> mkdf(5,3,data_gen_f=lambda r,c:randint(1,100))
  1665. # 2-level multiindex on rows with each label duplicated
  1666. # twice on first level, default names on both axis, single
  1667. # index on both axis
  1668. >> a=makeCustomDataframe(5,3,r_idx_nlevels=2,r_ndupe_l=[2])
  1669. # DatetimeIndex on row, index with unicode labels on columns
  1670. # no names on either axis
  1671. >> a=makeCustomDataframe(5,3,c_idx_names=False,r_idx_names=False,
  1672. r_idx_type="dt",c_idx_type="u")
  1673. # 4-level multindex on rows with names provided, 2-level multindex
  1674. # on columns with default labels and default names.
  1675. >> a=makeCustomDataframe(5,3,r_idx_nlevels=4,
  1676. r_idx_names=["FEE","FI","FO","FAM"],
  1677. c_idx_nlevels=2)
  1678. >> a=mkdf(5,3,r_idx_nlevels=2,c_idx_nlevels=4)
  1679. """
  1680. assert c_idx_nlevels > 0
  1681. assert r_idx_nlevels > 0
  1682. assert r_idx_type is None or (
  1683. r_idx_type in ("i", "f", "s", "u", "dt", "p", "td") and r_idx_nlevels == 1
  1684. )
  1685. assert c_idx_type is None or (
  1686. c_idx_type in ("i", "f", "s", "u", "dt", "p", "td") and c_idx_nlevels == 1
  1687. )
  1688. columns = makeCustomIndex(
  1689. ncols,
  1690. nlevels=c_idx_nlevels,
  1691. prefix="C",
  1692. names=c_idx_names,
  1693. ndupe_l=c_ndupe_l,
  1694. idx_type=c_idx_type,
  1695. )
  1696. index = makeCustomIndex(
  1697. nrows,
  1698. nlevels=r_idx_nlevels,
  1699. prefix="R",
  1700. names=r_idx_names,
  1701. ndupe_l=r_ndupe_l,
  1702. idx_type=r_idx_type,
  1703. )
  1704. # by default, generate data based on location
  1705. if data_gen_f is None:
  1706. data_gen_f = lambda r, c: f"R{r}C{c}"
  1707. data = [[data_gen_f(r, c) for c in range(ncols)] for r in range(nrows)]
  1708. return DataFrame(data, index, columns, dtype=dtype)
  1709. def _create_missing_idx(nrows, ncols, density, random_state=None):
  1710. if random_state is None:
  1711. random_state = np.random
  1712. else:
  1713. random_state = np.random.RandomState(random_state)
  1714. # below is cribbed from scipy.sparse
  1715. size = int(np.round((1 - density) * nrows * ncols))
  1716. # generate a few more to ensure unique values
  1717. min_rows = 5
  1718. fac = 1.02
  1719. extra_size = min(size + min_rows, fac * size)
  1720. def _gen_unique_rand(rng, _extra_size):
  1721. ind = rng.rand(int(_extra_size))
  1722. return np.unique(np.floor(ind * nrows * ncols))[:size]
  1723. ind = _gen_unique_rand(random_state, extra_size)
  1724. while ind.size < size:
  1725. extra_size *= 1.05
  1726. ind = _gen_unique_rand(random_state, extra_size)
  1727. j = np.floor(ind * 1.0 / nrows).astype(int)
  1728. i = (ind - j * nrows).astype(int)
  1729. return i.tolist(), j.tolist()
  1730. def makeMissingCustomDataframe(
  1731. nrows,
  1732. ncols,
  1733. density=0.9,
  1734. random_state=None,
  1735. c_idx_names=True,
  1736. r_idx_names=True,
  1737. c_idx_nlevels=1,
  1738. r_idx_nlevels=1,
  1739. data_gen_f=None,
  1740. c_ndupe_l=None,
  1741. r_ndupe_l=None,
  1742. dtype=None,
  1743. c_idx_type=None,
  1744. r_idx_type=None,
  1745. ):
  1746. """
  1747. Parameters
  1748. ----------
  1749. Density : float, optional
  1750. Float in (0, 1) that gives the percentage of non-missing numbers in
  1751. the DataFrame.
  1752. random_state : {np.random.RandomState, int}, optional
  1753. Random number generator or random seed.
  1754. See makeCustomDataframe for descriptions of the rest of the parameters.
  1755. """
  1756. df = makeCustomDataframe(
  1757. nrows,
  1758. ncols,
  1759. c_idx_names=c_idx_names,
  1760. r_idx_names=r_idx_names,
  1761. c_idx_nlevels=c_idx_nlevels,
  1762. r_idx_nlevels=r_idx_nlevels,
  1763. data_gen_f=data_gen_f,
  1764. c_ndupe_l=c_ndupe_l,
  1765. r_ndupe_l=r_ndupe_l,
  1766. dtype=dtype,
  1767. c_idx_type=c_idx_type,
  1768. r_idx_type=r_idx_type,
  1769. )
  1770. i, j = _create_missing_idx(nrows, ncols, density, random_state)
  1771. df.values[i, j] = np.nan
  1772. return df
  1773. def makeMissingDataframe(density=0.9, random_state=None):
  1774. df = makeDataFrame()
  1775. i, j = _create_missing_idx(*df.shape, density=density, random_state=random_state)
  1776. df.values[i, j] = np.nan
  1777. return df
  1778. def optional_args(decorator):
  1779. """allows a decorator to take optional positional and keyword arguments.
  1780. Assumes that taking a single, callable, positional argument means that
  1781. it is decorating a function, i.e. something like this::
  1782. @my_decorator
  1783. def function(): pass
  1784. Calls decorator with decorator(f, *args, **kwargs)"""
  1785. @wraps(decorator)
  1786. def wrapper(*args, **kwargs):
  1787. def dec(f):
  1788. return decorator(f, *args, **kwargs)
  1789. is_decorating = not kwargs and len(args) == 1 and callable(args[0])
  1790. if is_decorating:
  1791. f = args[0]
  1792. args = []
  1793. return dec(f)
  1794. else:
  1795. return dec
  1796. return wrapper
  1797. # skip tests on exceptions with this message
  1798. _network_error_messages = (
  1799. # 'urlopen error timed out',
  1800. # 'timeout: timed out',
  1801. # 'socket.timeout: timed out',
  1802. "timed out",
  1803. "Server Hangup",
  1804. "HTTP Error 503: Service Unavailable",
  1805. "502: Proxy Error",
  1806. "HTTP Error 502: internal error",
  1807. "HTTP Error 502",
  1808. "HTTP Error 503",
  1809. "HTTP Error 403",
  1810. "HTTP Error 400",
  1811. "Temporary failure in name resolution",
  1812. "Name or service not known",
  1813. "Connection refused",
  1814. "certificate verify",
  1815. )
  1816. # or this e.errno/e.reason.errno
  1817. _network_errno_vals = (
  1818. 101, # Network is unreachable
  1819. 111, # Connection refused
  1820. 110, # Connection timed out
  1821. 104, # Connection reset Error
  1822. 54, # Connection reset by peer
  1823. 60, # urllib.error.URLError: [Errno 60] Connection timed out
  1824. )
  1825. # Both of the above shouldn't mask real issues such as 404's
  1826. # or refused connections (changed DNS).
  1827. # But some tests (test_data yahoo) contact incredibly flakey
  1828. # servers.
  1829. # and conditionally raise on exception types in _get_default_network_errors
  1830. def _get_default_network_errors():
  1831. # Lazy import for http.client because it imports many things from the stdlib
  1832. import http.client
  1833. return (IOError, http.client.HTTPException, TimeoutError)
  1834. def can_connect(url, error_classes=None):
  1835. """Try to connect to the given url. True if succeeds, False if IOError
  1836. raised
  1837. Parameters
  1838. ----------
  1839. url : basestring
  1840. The URL to try to connect to
  1841. Returns
  1842. -------
  1843. connectable : bool
  1844. Return True if no IOError (unable to connect) or URLError (bad url) was
  1845. raised
  1846. """
  1847. if error_classes is None:
  1848. error_classes = _get_default_network_errors()
  1849. try:
  1850. with urlopen(url):
  1851. pass
  1852. except error_classes:
  1853. return False
  1854. else:
  1855. return True
  1856. @optional_args
  1857. def network(
  1858. t,
  1859. url="http://www.google.com",
  1860. raise_on_error=_RAISE_NETWORK_ERROR_DEFAULT,
  1861. check_before_test=False,
  1862. error_classes=None,
  1863. skip_errnos=_network_errno_vals,
  1864. _skip_on_messages=_network_error_messages,
  1865. ):
  1866. """
  1867. Label a test as requiring network connection and, if an error is
  1868. encountered, only raise if it does not find a network connection.
  1869. In comparison to ``network``, this assumes an added contract to your test:
  1870. you must assert that, under normal conditions, your test will ONLY fail if
  1871. it does not have network connectivity.
  1872. You can call this in 3 ways: as a standard decorator, with keyword
  1873. arguments, or with a positional argument that is the url to check.
  1874. Parameters
  1875. ----------
  1876. t : callable
  1877. The test requiring network connectivity.
  1878. url : path
  1879. The url to test via ``pandas.io.common.urlopen`` to check
  1880. for connectivity. Defaults to 'http://www.google.com'.
  1881. raise_on_error : bool
  1882. If True, never catches errors.
  1883. check_before_test : bool
  1884. If True, checks connectivity before running the test case.
  1885. error_classes : tuple or Exception
  1886. error classes to ignore. If not in ``error_classes``, raises the error.
  1887. defaults to IOError. Be careful about changing the error classes here.
  1888. skip_errnos : iterable of int
  1889. Any exception that has .errno or .reason.erno set to one
  1890. of these values will be skipped with an appropriate
  1891. message.
  1892. _skip_on_messages: iterable of string
  1893. any exception e for which one of the strings is
  1894. a substring of str(e) will be skipped with an appropriate
  1895. message. Intended to suppress errors where an errno isn't available.
  1896. Notes
  1897. -----
  1898. * ``raise_on_error`` supercedes ``check_before_test``
  1899. Returns
  1900. -------
  1901. t : callable
  1902. The decorated test ``t``, with checks for connectivity errors.
  1903. Example
  1904. -------
  1905. Tests decorated with @network will fail if it's possible to make a network
  1906. connection to another URL (defaults to google.com)::
  1907. >>> from pandas._testing import network
  1908. >>> from pandas.io.common import urlopen
  1909. >>> @network
  1910. ... def test_network():
  1911. ... with urlopen("rabbit://bonanza.com"):
  1912. ... pass
  1913. Traceback
  1914. ...
  1915. URLError: <urlopen error unknown url type: rabit>
  1916. You can specify alternative URLs::
  1917. >>> @network("http://www.yahoo.com")
  1918. ... def test_something_with_yahoo():
  1919. ... raise IOError("Failure Message")
  1920. >>> test_something_with_yahoo()
  1921. Traceback (most recent call last):
  1922. ...
  1923. IOError: Failure Message
  1924. If you set check_before_test, it will check the url first and not run the
  1925. test on failure::
  1926. >>> @network("failing://url.blaher", check_before_test=True)
  1927. ... def test_something():
  1928. ... print("I ran!")
  1929. ... raise ValueError("Failure")
  1930. >>> test_something()
  1931. Traceback (most recent call last):
  1932. ...
  1933. Errors not related to networking will always be raised.
  1934. """
  1935. from pytest import skip
  1936. if error_classes is None:
  1937. error_classes = _get_default_network_errors()
  1938. t.network = True
  1939. @wraps(t)
  1940. def wrapper(*args, **kwargs):
  1941. if check_before_test and not raise_on_error:
  1942. if not can_connect(url, error_classes):
  1943. skip()
  1944. try:
  1945. return t(*args, **kwargs)
  1946. except Exception as err:
  1947. errno = getattr(err, "errno", None)
  1948. if not errno and hasattr(errno, "reason"):
  1949. errno = getattr(err.reason, "errno", None)
  1950. if errno in skip_errnos:
  1951. skip(f"Skipping test due to known errno and error {err}")
  1952. e_str = str(err)
  1953. if any(m.lower() in e_str.lower() for m in _skip_on_messages):
  1954. skip(
  1955. f"Skipping test because exception message is known and error {err}"
  1956. )
  1957. if not isinstance(err, error_classes):
  1958. raise
  1959. if raise_on_error or can_connect(url, error_classes):
  1960. raise
  1961. else:
  1962. skip(f"Skipping test due to lack of connectivity and error {err}")
  1963. return wrapper
  1964. with_connectivity_check = network
  1965. @contextmanager
  1966. def assert_produces_warning(
  1967. expected_warning=Warning,
  1968. filter_level="always",
  1969. clear=None,
  1970. check_stacklevel=True,
  1971. raise_on_extra_warnings=True,
  1972. ):
  1973. """
  1974. Context manager for running code expected to either raise a specific
  1975. warning, or not raise any warnings. Verifies that the code raises the
  1976. expected warning, and that it does not raise any other unexpected
  1977. warnings. It is basically a wrapper around ``warnings.catch_warnings``.
  1978. Parameters
  1979. ----------
  1980. expected_warning : {Warning, False, None}, default Warning
  1981. The type of Exception raised. ``exception.Warning`` is the base
  1982. class for all warnings. To check that no warning is returned,
  1983. specify ``False`` or ``None``.
  1984. filter_level : str or None, default "always"
  1985. Specifies whether warnings are ignored, displayed, or turned
  1986. into errors.
  1987. Valid values are:
  1988. * "error" - turns matching warnings into exceptions
  1989. * "ignore" - discard the warning
  1990. * "always" - always emit a warning
  1991. * "default" - print the warning the first time it is generated
  1992. from each location
  1993. * "module" - print the warning the first time it is generated
  1994. from each module
  1995. * "once" - print the warning the first time it is generated
  1996. clear : str, default None
  1997. If not ``None`` then remove any previously raised warnings from
  1998. the ``__warningsregistry__`` to ensure that no warning messages are
  1999. suppressed by this context manager. If ``None`` is specified,
  2000. the ``__warningsregistry__`` keeps track of which warnings have been
  2001. shown, and does not show them again.
  2002. check_stacklevel : bool, default True
  2003. If True, displays the line that called the function containing
  2004. the warning to show were the function is called. Otherwise, the
  2005. line that implements the function is displayed.
  2006. raise_on_extra_warnings : bool, default True
  2007. Whether extra warnings not of the type `expected_warning` should
  2008. cause the test to fail.
  2009. Examples
  2010. --------
  2011. >>> import warnings
  2012. >>> with assert_produces_warning():
  2013. ... warnings.warn(UserWarning())
  2014. ...
  2015. >>> with assert_produces_warning(False):
  2016. ... warnings.warn(RuntimeWarning())
  2017. ...
  2018. Traceback (most recent call last):
  2019. ...
  2020. AssertionError: Caused unexpected warning(s): ['RuntimeWarning'].
  2021. >>> with assert_produces_warning(UserWarning):
  2022. ... warnings.warn(RuntimeWarning())
  2023. Traceback (most recent call last):
  2024. ...
  2025. AssertionError: Did not see expected warning of class 'UserWarning'.
  2026. ..warn:: This is *not* thread-safe.
  2027. """
  2028. __tracebackhide__ = True
  2029. with warnings.catch_warnings(record=True) as w:
  2030. if clear is not None:
  2031. # make sure that we are clearing these warnings
  2032. # if they have happened before
  2033. # to guarantee that we will catch them
  2034. if not is_list_like(clear):
  2035. clear = [clear]
  2036. for m in clear:
  2037. try:
  2038. m.__warningregistry__.clear()
  2039. except AttributeError:
  2040. # module may not have __warningregistry__
  2041. pass
  2042. saw_warning = False
  2043. warnings.simplefilter(filter_level)
  2044. yield w
  2045. extra_warnings = []
  2046. for actual_warning in w:
  2047. if expected_warning and issubclass(
  2048. actual_warning.category, expected_warning
  2049. ):
  2050. saw_warning = True
  2051. if check_stacklevel and issubclass(
  2052. actual_warning.category, (FutureWarning, DeprecationWarning)
  2053. ):
  2054. from inspect import getframeinfo, stack
  2055. caller = getframeinfo(stack()[2][0])
  2056. msg = (
  2057. "Warning not set with correct stacklevel. "
  2058. f"File where warning is raised: {actual_warning.filename} != "
  2059. f"{caller.filename}. Warning message: {actual_warning.message}"
  2060. )
  2061. assert actual_warning.filename == caller.filename, msg
  2062. else:
  2063. extra_warnings.append(
  2064. (
  2065. actual_warning.category.__name__,
  2066. actual_warning.message,
  2067. actual_warning.filename,
  2068. actual_warning.lineno,
  2069. )
  2070. )
  2071. if expected_warning:
  2072. msg = (
  2073. f"Did not see expected warning of class "
  2074. f"{repr(expected_warning.__name__)}"
  2075. )
  2076. assert saw_warning, msg
  2077. if raise_on_extra_warnings and extra_warnings:
  2078. raise AssertionError(
  2079. f"Caused unexpected warning(s): {repr(extra_warnings)}"
  2080. )
  2081. class RNGContext:
  2082. """
  2083. Context manager to set the numpy random number generator speed. Returns
  2084. to the original value upon exiting the context manager.
  2085. Parameters
  2086. ----------
  2087. seed : int
  2088. Seed for numpy.random.seed
  2089. Examples
  2090. --------
  2091. with RNGContext(42):
  2092. np.random.randn()
  2093. """
  2094. def __init__(self, seed):
  2095. self.seed = seed
  2096. def __enter__(self):
  2097. self.start_state = np.random.get_state()
  2098. np.random.seed(self.seed)
  2099. def __exit__(self, exc_type, exc_value, traceback):
  2100. np.random.set_state(self.start_state)
  2101. @contextmanager
  2102. def with_csv_dialect(name, **kwargs):
  2103. """
  2104. Context manager to temporarily register a CSV dialect for parsing CSV.
  2105. Parameters
  2106. ----------
  2107. name : str
  2108. The name of the dialect.
  2109. kwargs : mapping
  2110. The parameters for the dialect.
  2111. Raises
  2112. ------
  2113. ValueError : the name of the dialect conflicts with a builtin one.
  2114. See Also
  2115. --------
  2116. csv : Python's CSV library.
  2117. """
  2118. import csv
  2119. _BUILTIN_DIALECTS = {"excel", "excel-tab", "unix"}
  2120. if name in _BUILTIN_DIALECTS:
  2121. raise ValueError("Cannot override builtin dialect.")
  2122. csv.register_dialect(name, **kwargs)
  2123. yield
  2124. csv.unregister_dialect(name)
  2125. @contextmanager
  2126. def use_numexpr(use, min_elements=None):
  2127. from pandas.core.computation import expressions as expr
  2128. if min_elements is None:
  2129. min_elements = expr._MIN_ELEMENTS
  2130. olduse = expr._USE_NUMEXPR
  2131. oldmin = expr._MIN_ELEMENTS
  2132. expr.set_use_numexpr(use)
  2133. expr._MIN_ELEMENTS = min_elements
  2134. yield
  2135. expr._MIN_ELEMENTS = oldmin
  2136. expr.set_use_numexpr(olduse)
  2137. def test_parallel(num_threads=2, kwargs_list=None):
  2138. """Decorator to run the same function multiple times in parallel.
  2139. Parameters
  2140. ----------
  2141. num_threads : int, optional
  2142. The number of times the function is run in parallel.
  2143. kwargs_list : list of dicts, optional
  2144. The list of kwargs to update original
  2145. function kwargs on different threads.
  2146. Notes
  2147. -----
  2148. This decorator does not pass the return value of the decorated function.
  2149. Original from scikit-image:
  2150. https://github.com/scikit-image/scikit-image/pull/1519
  2151. """
  2152. assert num_threads > 0
  2153. has_kwargs_list = kwargs_list is not None
  2154. if has_kwargs_list:
  2155. assert len(kwargs_list) == num_threads
  2156. import threading
  2157. def wrapper(func):
  2158. @wraps(func)
  2159. def inner(*args, **kwargs):
  2160. if has_kwargs_list:
  2161. update_kwargs = lambda i: dict(kwargs, **kwargs_list[i])
  2162. else:
  2163. update_kwargs = lambda i: kwargs
  2164. threads = []
  2165. for i in range(num_threads):
  2166. updated_kwargs = update_kwargs(i)
  2167. thread = threading.Thread(target=func, args=args, kwargs=updated_kwargs)
  2168. threads.append(thread)
  2169. for thread in threads:
  2170. thread.start()
  2171. for thread in threads:
  2172. thread.join()
  2173. return inner
  2174. return wrapper
  2175. class SubclassedSeries(Series):
  2176. _metadata = ["testattr", "name"]
  2177. @property
  2178. def _constructor(self):
  2179. return SubclassedSeries
  2180. @property
  2181. def _constructor_expanddim(self):
  2182. return SubclassedDataFrame
  2183. class SubclassedDataFrame(DataFrame):
  2184. _metadata = ["testattr"]
  2185. @property
  2186. def _constructor(self):
  2187. return SubclassedDataFrame
  2188. @property
  2189. def _constructor_sliced(self):
  2190. return SubclassedSeries
  2191. class SubclassedCategorical(Categorical):
  2192. @property
  2193. def _constructor(self):
  2194. return SubclassedCategorical
  2195. @contextmanager
  2196. def set_timezone(tz: str):
  2197. """
  2198. Context manager for temporarily setting a timezone.
  2199. Parameters
  2200. ----------
  2201. tz : str
  2202. A string representing a valid timezone.
  2203. Examples
  2204. --------
  2205. >>> from datetime import datetime
  2206. >>> from dateutil.tz import tzlocal
  2207. >>> tzlocal().tzname(datetime.now())
  2208. 'IST'
  2209. >>> with set_timezone('US/Eastern'):
  2210. ... tzlocal().tzname(datetime.now())
  2211. ...
  2212. 'EDT'
  2213. """
  2214. import os
  2215. import time
  2216. def setTZ(tz):
  2217. if tz is None:
  2218. try:
  2219. del os.environ["TZ"]
  2220. except KeyError:
  2221. pass
  2222. else:
  2223. os.environ["TZ"] = tz
  2224. time.tzset()
  2225. orig_tz = os.environ.get("TZ")
  2226. setTZ(tz)
  2227. try:
  2228. yield
  2229. finally:
  2230. setTZ(orig_tz)
  2231. def _make_skipna_wrapper(alternative, skipna_alternative=None):
  2232. """
  2233. Create a function for calling on an array.
  2234. Parameters
  2235. ----------
  2236. alternative : function
  2237. The function to be called on the array with no NaNs.
  2238. Only used when 'skipna_alternative' is None.
  2239. skipna_alternative : function
  2240. The function to be called on the original array
  2241. Returns
  2242. -------
  2243. function
  2244. """
  2245. if skipna_alternative:
  2246. def skipna_wrapper(x):
  2247. return skipna_alternative(x.values)
  2248. else:
  2249. def skipna_wrapper(x):
  2250. nona = x.dropna()
  2251. if len(nona) == 0:
  2252. return np.nan
  2253. return alternative(nona)
  2254. return skipna_wrapper
  2255. def convert_rows_list_to_csv_str(rows_list: List[str]):
  2256. """
  2257. Convert list of CSV rows to single CSV-formatted string for current OS.
  2258. This method is used for creating expected value of to_csv() method.
  2259. Parameters
  2260. ----------
  2261. rows_list : List[str]
  2262. Each element represents the row of csv.
  2263. Returns
  2264. -------
  2265. str
  2266. Expected output of to_csv() in current OS.
  2267. """
  2268. sep = os.linesep
  2269. expected = sep.join(rows_list) + sep
  2270. return expected