test_subclass.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738
  1. import numpy as np
  2. import pytest
  3. import pandas.util._test_decorators as td
  4. import pandas as pd
  5. from pandas import (
  6. DataFrame,
  7. Index,
  8. MultiIndex,
  9. Series,
  10. )
  11. import pandas._testing as tm
  12. @pytest.fixture()
  13. def gpd_style_subclass_df():
  14. class SubclassedDataFrame(DataFrame):
  15. @property
  16. def _constructor(self):
  17. return SubclassedDataFrame
  18. return SubclassedDataFrame({"a": [1, 2, 3]})
  19. class TestDataFrameSubclassing:
  20. def test_frame_subclassing_and_slicing(self):
  21. # Subclass frame and ensure it returns the right class on slicing it
  22. # In reference to PR 9632
  23. class CustomSeries(Series):
  24. @property
  25. def _constructor(self):
  26. return CustomSeries
  27. def custom_series_function(self):
  28. return "OK"
  29. class CustomDataFrame(DataFrame):
  30. """
  31. Subclasses pandas DF, fills DF with simulation results, adds some
  32. custom plotting functions.
  33. """
  34. def __init__(self, *args, **kw) -> None:
  35. super().__init__(*args, **kw)
  36. @property
  37. def _constructor(self):
  38. return CustomDataFrame
  39. _constructor_sliced = CustomSeries
  40. def custom_frame_function(self):
  41. return "OK"
  42. data = {"col1": range(10), "col2": range(10)}
  43. cdf = CustomDataFrame(data)
  44. # Did we get back our own DF class?
  45. assert isinstance(cdf, CustomDataFrame)
  46. # Do we get back our own Series class after selecting a column?
  47. cdf_series = cdf.col1
  48. assert isinstance(cdf_series, CustomSeries)
  49. assert cdf_series.custom_series_function() == "OK"
  50. # Do we get back our own DF class after slicing row-wise?
  51. cdf_rows = cdf[1:5]
  52. assert isinstance(cdf_rows, CustomDataFrame)
  53. assert cdf_rows.custom_frame_function() == "OK"
  54. # Make sure sliced part of multi-index frame is custom class
  55. mcol = MultiIndex.from_tuples([("A", "A"), ("A", "B")])
  56. cdf_multi = CustomDataFrame([[0, 1], [2, 3]], columns=mcol)
  57. assert isinstance(cdf_multi["A"], CustomDataFrame)
  58. mcol = MultiIndex.from_tuples([("A", ""), ("B", "")])
  59. cdf_multi2 = CustomDataFrame([[0, 1], [2, 3]], columns=mcol)
  60. assert isinstance(cdf_multi2["A"], CustomSeries)
  61. def test_dataframe_metadata(self):
  62. df = tm.SubclassedDataFrame(
  63. {"X": [1, 2, 3], "Y": [1, 2, 3]}, index=["a", "b", "c"]
  64. )
  65. df.testattr = "XXX"
  66. assert df.testattr == "XXX"
  67. assert df[["X"]].testattr == "XXX"
  68. assert df.loc[["a", "b"], :].testattr == "XXX"
  69. assert df.iloc[[0, 1], :].testattr == "XXX"
  70. # see gh-9776
  71. assert df.iloc[0:1, :].testattr == "XXX"
  72. # see gh-10553
  73. unpickled = tm.round_trip_pickle(df)
  74. tm.assert_frame_equal(df, unpickled)
  75. assert df._metadata == unpickled._metadata
  76. assert df.testattr == unpickled.testattr
  77. def test_indexing_sliced(self):
  78. # GH 11559
  79. df = tm.SubclassedDataFrame(
  80. {"X": [1, 2, 3], "Y": [4, 5, 6], "Z": [7, 8, 9]}, index=["a", "b", "c"]
  81. )
  82. res = df.loc[:, "X"]
  83. exp = tm.SubclassedSeries([1, 2, 3], index=list("abc"), name="X")
  84. tm.assert_series_equal(res, exp)
  85. assert isinstance(res, tm.SubclassedSeries)
  86. res = df.iloc[:, 1]
  87. exp = tm.SubclassedSeries([4, 5, 6], index=list("abc"), name="Y")
  88. tm.assert_series_equal(res, exp)
  89. assert isinstance(res, tm.SubclassedSeries)
  90. res = df.loc[:, "Z"]
  91. exp = tm.SubclassedSeries([7, 8, 9], index=list("abc"), name="Z")
  92. tm.assert_series_equal(res, exp)
  93. assert isinstance(res, tm.SubclassedSeries)
  94. res = df.loc["a", :]
  95. exp = tm.SubclassedSeries([1, 4, 7], index=list("XYZ"), name="a")
  96. tm.assert_series_equal(res, exp)
  97. assert isinstance(res, tm.SubclassedSeries)
  98. res = df.iloc[1, :]
  99. exp = tm.SubclassedSeries([2, 5, 8], index=list("XYZ"), name="b")
  100. tm.assert_series_equal(res, exp)
  101. assert isinstance(res, tm.SubclassedSeries)
  102. res = df.loc["c", :]
  103. exp = tm.SubclassedSeries([3, 6, 9], index=list("XYZ"), name="c")
  104. tm.assert_series_equal(res, exp)
  105. assert isinstance(res, tm.SubclassedSeries)
  106. def test_subclass_attr_err_propagation(self):
  107. # GH 11808
  108. class A(DataFrame):
  109. @property
  110. def nonexistence(self):
  111. return self.i_dont_exist
  112. with pytest.raises(AttributeError, match=".*i_dont_exist.*"):
  113. A().nonexistence
  114. def test_subclass_align(self):
  115. # GH 12983
  116. df1 = tm.SubclassedDataFrame(
  117. {"a": [1, 3, 5], "b": [1, 3, 5]}, index=list("ACE")
  118. )
  119. df2 = tm.SubclassedDataFrame(
  120. {"c": [1, 2, 4], "d": [1, 2, 4]}, index=list("ABD")
  121. )
  122. res1, res2 = df1.align(df2, axis=0)
  123. exp1 = tm.SubclassedDataFrame(
  124. {"a": [1, np.nan, 3, np.nan, 5], "b": [1, np.nan, 3, np.nan, 5]},
  125. index=list("ABCDE"),
  126. )
  127. exp2 = tm.SubclassedDataFrame(
  128. {"c": [1, 2, np.nan, 4, np.nan], "d": [1, 2, np.nan, 4, np.nan]},
  129. index=list("ABCDE"),
  130. )
  131. assert isinstance(res1, tm.SubclassedDataFrame)
  132. tm.assert_frame_equal(res1, exp1)
  133. assert isinstance(res2, tm.SubclassedDataFrame)
  134. tm.assert_frame_equal(res2, exp2)
  135. res1, res2 = df1.a.align(df2.c)
  136. assert isinstance(res1, tm.SubclassedSeries)
  137. tm.assert_series_equal(res1, exp1.a)
  138. assert isinstance(res2, tm.SubclassedSeries)
  139. tm.assert_series_equal(res2, exp2.c)
  140. def test_subclass_align_combinations(self):
  141. # GH 12983
  142. df = tm.SubclassedDataFrame({"a": [1, 3, 5], "b": [1, 3, 5]}, index=list("ACE"))
  143. s = tm.SubclassedSeries([1, 2, 4], index=list("ABD"), name="x")
  144. # frame + series
  145. res1, res2 = df.align(s, axis=0)
  146. exp1 = tm.SubclassedDataFrame(
  147. {"a": [1, np.nan, 3, np.nan, 5], "b": [1, np.nan, 3, np.nan, 5]},
  148. index=list("ABCDE"),
  149. )
  150. # name is lost when
  151. exp2 = tm.SubclassedSeries(
  152. [1, 2, np.nan, 4, np.nan], index=list("ABCDE"), name="x"
  153. )
  154. assert isinstance(res1, tm.SubclassedDataFrame)
  155. tm.assert_frame_equal(res1, exp1)
  156. assert isinstance(res2, tm.SubclassedSeries)
  157. tm.assert_series_equal(res2, exp2)
  158. # series + frame
  159. res1, res2 = s.align(df)
  160. assert isinstance(res1, tm.SubclassedSeries)
  161. tm.assert_series_equal(res1, exp2)
  162. assert isinstance(res2, tm.SubclassedDataFrame)
  163. tm.assert_frame_equal(res2, exp1)
  164. def test_subclass_iterrows(self):
  165. # GH 13977
  166. df = tm.SubclassedDataFrame({"a": [1]})
  167. for i, row in df.iterrows():
  168. assert isinstance(row, tm.SubclassedSeries)
  169. tm.assert_series_equal(row, df.loc[i])
  170. def test_subclass_stack(self):
  171. # GH 15564
  172. df = tm.SubclassedDataFrame(
  173. [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
  174. index=["a", "b", "c"],
  175. columns=["X", "Y", "Z"],
  176. )
  177. res = df.stack()
  178. exp = tm.SubclassedSeries(
  179. [1, 2, 3, 4, 5, 6, 7, 8, 9], index=[list("aaabbbccc"), list("XYZXYZXYZ")]
  180. )
  181. tm.assert_series_equal(res, exp)
  182. def test_subclass_stack_multi(self):
  183. # GH 15564
  184. df = tm.SubclassedDataFrame(
  185. [[10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33], [40, 41, 42, 43]],
  186. index=MultiIndex.from_tuples(
  187. list(zip(list("AABB"), list("cdcd"))), names=["aaa", "ccc"]
  188. ),
  189. columns=MultiIndex.from_tuples(
  190. list(zip(list("WWXX"), list("yzyz"))), names=["www", "yyy"]
  191. ),
  192. )
  193. exp = tm.SubclassedDataFrame(
  194. [
  195. [10, 12],
  196. [11, 13],
  197. [20, 22],
  198. [21, 23],
  199. [30, 32],
  200. [31, 33],
  201. [40, 42],
  202. [41, 43],
  203. ],
  204. index=MultiIndex.from_tuples(
  205. list(zip(list("AAAABBBB"), list("ccddccdd"), list("yzyzyzyz"))),
  206. names=["aaa", "ccc", "yyy"],
  207. ),
  208. columns=Index(["W", "X"], name="www"),
  209. )
  210. res = df.stack()
  211. tm.assert_frame_equal(res, exp)
  212. res = df.stack("yyy")
  213. tm.assert_frame_equal(res, exp)
  214. exp = tm.SubclassedDataFrame(
  215. [
  216. [10, 11],
  217. [12, 13],
  218. [20, 21],
  219. [22, 23],
  220. [30, 31],
  221. [32, 33],
  222. [40, 41],
  223. [42, 43],
  224. ],
  225. index=MultiIndex.from_tuples(
  226. list(zip(list("AAAABBBB"), list("ccddccdd"), list("WXWXWXWX"))),
  227. names=["aaa", "ccc", "www"],
  228. ),
  229. columns=Index(["y", "z"], name="yyy"),
  230. )
  231. res = df.stack("www")
  232. tm.assert_frame_equal(res, exp)
  233. def test_subclass_stack_multi_mixed(self):
  234. # GH 15564
  235. df = tm.SubclassedDataFrame(
  236. [
  237. [10, 11, 12.0, 13.0],
  238. [20, 21, 22.0, 23.0],
  239. [30, 31, 32.0, 33.0],
  240. [40, 41, 42.0, 43.0],
  241. ],
  242. index=MultiIndex.from_tuples(
  243. list(zip(list("AABB"), list("cdcd"))), names=["aaa", "ccc"]
  244. ),
  245. columns=MultiIndex.from_tuples(
  246. list(zip(list("WWXX"), list("yzyz"))), names=["www", "yyy"]
  247. ),
  248. )
  249. exp = tm.SubclassedDataFrame(
  250. [
  251. [10, 12.0],
  252. [11, 13.0],
  253. [20, 22.0],
  254. [21, 23.0],
  255. [30, 32.0],
  256. [31, 33.0],
  257. [40, 42.0],
  258. [41, 43.0],
  259. ],
  260. index=MultiIndex.from_tuples(
  261. list(zip(list("AAAABBBB"), list("ccddccdd"), list("yzyzyzyz"))),
  262. names=["aaa", "ccc", "yyy"],
  263. ),
  264. columns=Index(["W", "X"], name="www"),
  265. )
  266. res = df.stack()
  267. tm.assert_frame_equal(res, exp)
  268. res = df.stack("yyy")
  269. tm.assert_frame_equal(res, exp)
  270. exp = tm.SubclassedDataFrame(
  271. [
  272. [10.0, 11.0],
  273. [12.0, 13.0],
  274. [20.0, 21.0],
  275. [22.0, 23.0],
  276. [30.0, 31.0],
  277. [32.0, 33.0],
  278. [40.0, 41.0],
  279. [42.0, 43.0],
  280. ],
  281. index=MultiIndex.from_tuples(
  282. list(zip(list("AAAABBBB"), list("ccddccdd"), list("WXWXWXWX"))),
  283. names=["aaa", "ccc", "www"],
  284. ),
  285. columns=Index(["y", "z"], name="yyy"),
  286. )
  287. res = df.stack("www")
  288. tm.assert_frame_equal(res, exp)
  289. def test_subclass_unstack(self):
  290. # GH 15564
  291. df = tm.SubclassedDataFrame(
  292. [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
  293. index=["a", "b", "c"],
  294. columns=["X", "Y", "Z"],
  295. )
  296. res = df.unstack()
  297. exp = tm.SubclassedSeries(
  298. [1, 4, 7, 2, 5, 8, 3, 6, 9], index=[list("XXXYYYZZZ"), list("abcabcabc")]
  299. )
  300. tm.assert_series_equal(res, exp)
  301. def test_subclass_unstack_multi(self):
  302. # GH 15564
  303. df = tm.SubclassedDataFrame(
  304. [[10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33], [40, 41, 42, 43]],
  305. index=MultiIndex.from_tuples(
  306. list(zip(list("AABB"), list("cdcd"))), names=["aaa", "ccc"]
  307. ),
  308. columns=MultiIndex.from_tuples(
  309. list(zip(list("WWXX"), list("yzyz"))), names=["www", "yyy"]
  310. ),
  311. )
  312. exp = tm.SubclassedDataFrame(
  313. [[10, 20, 11, 21, 12, 22, 13, 23], [30, 40, 31, 41, 32, 42, 33, 43]],
  314. index=Index(["A", "B"], name="aaa"),
  315. columns=MultiIndex.from_tuples(
  316. list(zip(list("WWWWXXXX"), list("yyzzyyzz"), list("cdcdcdcd"))),
  317. names=["www", "yyy", "ccc"],
  318. ),
  319. )
  320. res = df.unstack()
  321. tm.assert_frame_equal(res, exp)
  322. res = df.unstack("ccc")
  323. tm.assert_frame_equal(res, exp)
  324. exp = tm.SubclassedDataFrame(
  325. [[10, 30, 11, 31, 12, 32, 13, 33], [20, 40, 21, 41, 22, 42, 23, 43]],
  326. index=Index(["c", "d"], name="ccc"),
  327. columns=MultiIndex.from_tuples(
  328. list(zip(list("WWWWXXXX"), list("yyzzyyzz"), list("ABABABAB"))),
  329. names=["www", "yyy", "aaa"],
  330. ),
  331. )
  332. res = df.unstack("aaa")
  333. tm.assert_frame_equal(res, exp)
  334. def test_subclass_unstack_multi_mixed(self):
  335. # GH 15564
  336. df = tm.SubclassedDataFrame(
  337. [
  338. [10, 11, 12.0, 13.0],
  339. [20, 21, 22.0, 23.0],
  340. [30, 31, 32.0, 33.0],
  341. [40, 41, 42.0, 43.0],
  342. ],
  343. index=MultiIndex.from_tuples(
  344. list(zip(list("AABB"), list("cdcd"))), names=["aaa", "ccc"]
  345. ),
  346. columns=MultiIndex.from_tuples(
  347. list(zip(list("WWXX"), list("yzyz"))), names=["www", "yyy"]
  348. ),
  349. )
  350. exp = tm.SubclassedDataFrame(
  351. [
  352. [10, 20, 11, 21, 12.0, 22.0, 13.0, 23.0],
  353. [30, 40, 31, 41, 32.0, 42.0, 33.0, 43.0],
  354. ],
  355. index=Index(["A", "B"], name="aaa"),
  356. columns=MultiIndex.from_tuples(
  357. list(zip(list("WWWWXXXX"), list("yyzzyyzz"), list("cdcdcdcd"))),
  358. names=["www", "yyy", "ccc"],
  359. ),
  360. )
  361. res = df.unstack()
  362. tm.assert_frame_equal(res, exp)
  363. res = df.unstack("ccc")
  364. tm.assert_frame_equal(res, exp)
  365. exp = tm.SubclassedDataFrame(
  366. [
  367. [10, 30, 11, 31, 12.0, 32.0, 13.0, 33.0],
  368. [20, 40, 21, 41, 22.0, 42.0, 23.0, 43.0],
  369. ],
  370. index=Index(["c", "d"], name="ccc"),
  371. columns=MultiIndex.from_tuples(
  372. list(zip(list("WWWWXXXX"), list("yyzzyyzz"), list("ABABABAB"))),
  373. names=["www", "yyy", "aaa"],
  374. ),
  375. )
  376. res = df.unstack("aaa")
  377. tm.assert_frame_equal(res, exp)
  378. def test_subclass_pivot(self):
  379. # GH 15564
  380. df = tm.SubclassedDataFrame(
  381. {
  382. "index": ["A", "B", "C", "C", "B", "A"],
  383. "columns": ["One", "One", "One", "Two", "Two", "Two"],
  384. "values": [1.0, 2.0, 3.0, 3.0, 2.0, 1.0],
  385. }
  386. )
  387. pivoted = df.pivot(index="index", columns="columns", values="values")
  388. expected = tm.SubclassedDataFrame(
  389. {
  390. "One": {"A": 1.0, "B": 2.0, "C": 3.0},
  391. "Two": {"A": 1.0, "B": 2.0, "C": 3.0},
  392. }
  393. )
  394. expected.index.name, expected.columns.name = "index", "columns"
  395. tm.assert_frame_equal(pivoted, expected)
  396. def test_subclassed_melt(self):
  397. # GH 15564
  398. cheese = tm.SubclassedDataFrame(
  399. {
  400. "first": ["John", "Mary"],
  401. "last": ["Doe", "Bo"],
  402. "height": [5.5, 6.0],
  403. "weight": [130, 150],
  404. }
  405. )
  406. melted = pd.melt(cheese, id_vars=["first", "last"])
  407. expected = tm.SubclassedDataFrame(
  408. [
  409. ["John", "Doe", "height", 5.5],
  410. ["Mary", "Bo", "height", 6.0],
  411. ["John", "Doe", "weight", 130],
  412. ["Mary", "Bo", "weight", 150],
  413. ],
  414. columns=["first", "last", "variable", "value"],
  415. )
  416. tm.assert_frame_equal(melted, expected)
  417. def test_subclassed_wide_to_long(self):
  418. # GH 9762
  419. np.random.seed(123)
  420. x = np.random.randn(3)
  421. df = tm.SubclassedDataFrame(
  422. {
  423. "A1970": {0: "a", 1: "b", 2: "c"},
  424. "A1980": {0: "d", 1: "e", 2: "f"},
  425. "B1970": {0: 2.5, 1: 1.2, 2: 0.7},
  426. "B1980": {0: 3.2, 1: 1.3, 2: 0.1},
  427. "X": dict(zip(range(3), x)),
  428. }
  429. )
  430. df["id"] = df.index
  431. exp_data = {
  432. "X": x.tolist() + x.tolist(),
  433. "A": ["a", "b", "c", "d", "e", "f"],
  434. "B": [2.5, 1.2, 0.7, 3.2, 1.3, 0.1],
  435. "year": [1970, 1970, 1970, 1980, 1980, 1980],
  436. "id": [0, 1, 2, 0, 1, 2],
  437. }
  438. expected = tm.SubclassedDataFrame(exp_data)
  439. expected = expected.set_index(["id", "year"])[["X", "A", "B"]]
  440. long_frame = pd.wide_to_long(df, ["A", "B"], i="id", j="year")
  441. tm.assert_frame_equal(long_frame, expected)
  442. def test_subclassed_apply(self):
  443. # GH 19822
  444. def check_row_subclass(row):
  445. assert isinstance(row, tm.SubclassedSeries)
  446. def stretch(row):
  447. if row["variable"] == "height":
  448. row["value"] += 0.5
  449. return row
  450. df = tm.SubclassedDataFrame(
  451. [
  452. ["John", "Doe", "height", 5.5],
  453. ["Mary", "Bo", "height", 6.0],
  454. ["John", "Doe", "weight", 130],
  455. ["Mary", "Bo", "weight", 150],
  456. ],
  457. columns=["first", "last", "variable", "value"],
  458. )
  459. df.apply(lambda x: check_row_subclass(x))
  460. df.apply(lambda x: check_row_subclass(x), axis=1)
  461. expected = tm.SubclassedDataFrame(
  462. [
  463. ["John", "Doe", "height", 6.0],
  464. ["Mary", "Bo", "height", 6.5],
  465. ["John", "Doe", "weight", 130],
  466. ["Mary", "Bo", "weight", 150],
  467. ],
  468. columns=["first", "last", "variable", "value"],
  469. )
  470. result = df.apply(lambda x: stretch(x), axis=1)
  471. assert isinstance(result, tm.SubclassedDataFrame)
  472. tm.assert_frame_equal(result, expected)
  473. expected = tm.SubclassedDataFrame([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]])
  474. result = df.apply(lambda x: tm.SubclassedSeries([1, 2, 3]), axis=1)
  475. assert isinstance(result, tm.SubclassedDataFrame)
  476. tm.assert_frame_equal(result, expected)
  477. result = df.apply(lambda x: [1, 2, 3], axis=1, result_type="expand")
  478. assert isinstance(result, tm.SubclassedDataFrame)
  479. tm.assert_frame_equal(result, expected)
  480. expected = tm.SubclassedSeries([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]])
  481. result = df.apply(lambda x: [1, 2, 3], axis=1)
  482. assert not isinstance(result, tm.SubclassedDataFrame)
  483. tm.assert_series_equal(result, expected)
  484. def test_subclassed_reductions(self, all_reductions):
  485. # GH 25596
  486. df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
  487. result = getattr(df, all_reductions)()
  488. assert isinstance(result, tm.SubclassedSeries)
  489. def test_subclassed_count(self):
  490. df = tm.SubclassedDataFrame(
  491. {
  492. "Person": ["John", "Myla", "Lewis", "John", "Myla"],
  493. "Age": [24.0, np.nan, 21.0, 33, 26],
  494. "Single": [False, True, True, True, False],
  495. }
  496. )
  497. result = df.count()
  498. assert isinstance(result, tm.SubclassedSeries)
  499. df = tm.SubclassedDataFrame({"A": [1, 0, 3], "B": [0, 5, 6], "C": [7, 8, 0]})
  500. result = df.count()
  501. assert isinstance(result, tm.SubclassedSeries)
  502. df = tm.SubclassedDataFrame(
  503. [[10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33], [40, 41, 42, 43]],
  504. index=MultiIndex.from_tuples(
  505. list(zip(list("AABB"), list("cdcd"))), names=["aaa", "ccc"]
  506. ),
  507. columns=MultiIndex.from_tuples(
  508. list(zip(list("WWXX"), list("yzyz"))), names=["www", "yyy"]
  509. ),
  510. )
  511. result = df.count()
  512. assert isinstance(result, tm.SubclassedSeries)
  513. df = tm.SubclassedDataFrame()
  514. result = df.count()
  515. assert isinstance(result, tm.SubclassedSeries)
  516. def test_isin(self):
  517. df = tm.SubclassedDataFrame(
  518. {"num_legs": [2, 4], "num_wings": [2, 0]}, index=["falcon", "dog"]
  519. )
  520. result = df.isin([0, 2])
  521. assert isinstance(result, tm.SubclassedDataFrame)
  522. def test_duplicated(self):
  523. df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
  524. result = df.duplicated()
  525. assert isinstance(result, tm.SubclassedSeries)
  526. df = tm.SubclassedDataFrame()
  527. result = df.duplicated()
  528. assert isinstance(result, tm.SubclassedSeries)
  529. @pytest.mark.parametrize("idx_method", ["idxmax", "idxmin"])
  530. def test_idx(self, idx_method):
  531. df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
  532. result = getattr(df, idx_method)()
  533. assert isinstance(result, tm.SubclassedSeries)
  534. def test_dot(self):
  535. df = tm.SubclassedDataFrame([[0, 1, -2, -1], [1, 1, 1, 1]])
  536. s = tm.SubclassedSeries([1, 1, 2, 1])
  537. result = df.dot(s)
  538. assert isinstance(result, tm.SubclassedSeries)
  539. df = tm.SubclassedDataFrame([[0, 1, -2, -1], [1, 1, 1, 1]])
  540. s = tm.SubclassedDataFrame([1, 1, 2, 1])
  541. result = df.dot(s)
  542. assert isinstance(result, tm.SubclassedDataFrame)
  543. def test_memory_usage(self):
  544. df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
  545. result = df.memory_usage()
  546. assert isinstance(result, tm.SubclassedSeries)
  547. result = df.memory_usage(index=False)
  548. assert isinstance(result, tm.SubclassedSeries)
  549. @td.skip_if_no_scipy
  550. def test_corrwith(self):
  551. index = ["a", "b", "c", "d", "e"]
  552. columns = ["one", "two", "three", "four"]
  553. df1 = tm.SubclassedDataFrame(
  554. np.random.randn(5, 4), index=index, columns=columns
  555. )
  556. df2 = tm.SubclassedDataFrame(
  557. np.random.randn(4, 4), index=index[:4], columns=columns
  558. )
  559. correls = df1.corrwith(df2, axis=1, drop=True, method="kendall")
  560. assert isinstance(correls, (tm.SubclassedSeries))
  561. def test_asof(self):
  562. N = 3
  563. rng = pd.date_range("1/1/1990", periods=N, freq="53s")
  564. df = tm.SubclassedDataFrame(
  565. {
  566. "A": [np.nan, np.nan, np.nan],
  567. "B": [np.nan, np.nan, np.nan],
  568. "C": [np.nan, np.nan, np.nan],
  569. },
  570. index=rng,
  571. )
  572. result = df.asof(rng[-2:])
  573. assert isinstance(result, tm.SubclassedDataFrame)
  574. result = df.asof(rng[-2])
  575. assert isinstance(result, tm.SubclassedSeries)
  576. result = df.asof("1989-12-31")
  577. assert isinstance(result, tm.SubclassedSeries)
  578. def test_idxmin_preserves_subclass(self):
  579. # GH 28330
  580. df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
  581. result = df.idxmin()
  582. assert isinstance(result, tm.SubclassedSeries)
  583. def test_idxmax_preserves_subclass(self):
  584. # GH 28330
  585. df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
  586. result = df.idxmax()
  587. assert isinstance(result, tm.SubclassedSeries)
  588. def test_convert_dtypes_preserves_subclass(self, gpd_style_subclass_df):
  589. # GH 43668
  590. df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
  591. result = df.convert_dtypes()
  592. assert isinstance(result, tm.SubclassedDataFrame)
  593. result = gpd_style_subclass_df.convert_dtypes()
  594. assert isinstance(result, type(gpd_style_subclass_df))
  595. def test_astype_preserves_subclass(self):
  596. # GH#40810
  597. df = tm.SubclassedDataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
  598. result = df.astype({"A": np.int64, "B": np.int32, "C": np.float64})
  599. assert isinstance(result, tm.SubclassedDataFrame)
  600. def test_equals_subclass(self):
  601. # https://github.com/pandas-dev/pandas/pull/34402
  602. # allow subclass in both directions
  603. df1 = DataFrame({"a": [1, 2, 3]})
  604. df2 = tm.SubclassedDataFrame({"a": [1, 2, 3]})
  605. assert df1.equals(df2)
  606. assert df2.equals(df1)
  607. def test_replace_list_method(self):
  608. # https://github.com/pandas-dev/pandas/pull/46018
  609. df = tm.SubclassedDataFrame({"A": [0, 1, 2]})
  610. result = df.replace([1, 2], method="ffill")
  611. expected = tm.SubclassedDataFrame({"A": [0, 0, 0]})
  612. assert isinstance(result, tm.SubclassedDataFrame)
  613. tm.assert_frame_equal(result, expected)