test_base.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364
  1. import os
  2. import shutil
  3. import tempfile
  4. import warnings
  5. from functools import partial
  6. from pathlib import Path
  7. from pickle import dumps, loads
  8. import numpy as np
  9. import pytest
  10. from sklearn.datasets import (
  11. clear_data_home,
  12. get_data_home,
  13. load_breast_cancer,
  14. load_diabetes,
  15. load_digits,
  16. load_files,
  17. load_iris,
  18. load_linnerud,
  19. load_sample_image,
  20. load_sample_images,
  21. load_wine,
  22. )
  23. from sklearn.datasets._base import (
  24. load_csv_data,
  25. load_gzip_compressed_csv_data,
  26. )
  27. from sklearn.datasets.tests.test_common import check_as_frame
  28. from sklearn.preprocessing import scale
  29. from sklearn.utils import Bunch
  30. from sklearn.utils.fixes import _is_resource
  31. class _DummyPath:
  32. """Minimal class that implements the os.PathLike interface."""
  33. def __init__(self, path):
  34. self.path = path
  35. def __fspath__(self):
  36. return self.path
  37. def _remove_dir(path):
  38. if os.path.isdir(path):
  39. shutil.rmtree(path)
  40. @pytest.fixture(scope="module")
  41. def data_home(tmpdir_factory):
  42. tmp_file = str(tmpdir_factory.mktemp("scikit_learn_data_home_test"))
  43. yield tmp_file
  44. _remove_dir(tmp_file)
  45. @pytest.fixture(scope="module")
  46. def load_files_root(tmpdir_factory):
  47. tmp_file = str(tmpdir_factory.mktemp("scikit_learn_load_files_test"))
  48. yield tmp_file
  49. _remove_dir(tmp_file)
  50. @pytest.fixture
  51. def test_category_dir_1(load_files_root):
  52. test_category_dir1 = tempfile.mkdtemp(dir=load_files_root)
  53. sample_file = tempfile.NamedTemporaryFile(dir=test_category_dir1, delete=False)
  54. sample_file.write(b"Hello World!\n")
  55. sample_file.close()
  56. yield str(test_category_dir1)
  57. _remove_dir(test_category_dir1)
  58. @pytest.fixture
  59. def test_category_dir_2(load_files_root):
  60. test_category_dir2 = tempfile.mkdtemp(dir=load_files_root)
  61. yield str(test_category_dir2)
  62. _remove_dir(test_category_dir2)
  63. @pytest.mark.parametrize("path_container", [None, Path, _DummyPath])
  64. def test_data_home(path_container, data_home):
  65. # get_data_home will point to a pre-existing folder
  66. if path_container is not None:
  67. data_home = path_container(data_home)
  68. data_home = get_data_home(data_home=data_home)
  69. assert data_home == data_home
  70. assert os.path.exists(data_home)
  71. # clear_data_home will delete both the content and the folder it-self
  72. if path_container is not None:
  73. data_home = path_container(data_home)
  74. clear_data_home(data_home=data_home)
  75. assert not os.path.exists(data_home)
  76. # if the folder is missing it will be created again
  77. data_home = get_data_home(data_home=data_home)
  78. assert os.path.exists(data_home)
  79. def test_default_empty_load_files(load_files_root):
  80. res = load_files(load_files_root)
  81. assert len(res.filenames) == 0
  82. assert len(res.target_names) == 0
  83. assert res.DESCR is None
  84. def test_default_load_files(test_category_dir_1, test_category_dir_2, load_files_root):
  85. res = load_files(load_files_root)
  86. assert len(res.filenames) == 1
  87. assert len(res.target_names) == 2
  88. assert res.DESCR is None
  89. assert res.data == [b"Hello World!\n"]
  90. def test_load_files_w_categories_desc_and_encoding(
  91. test_category_dir_1, test_category_dir_2, load_files_root
  92. ):
  93. category = os.path.abspath(test_category_dir_1).split(os.sep).pop()
  94. res = load_files(
  95. load_files_root, description="test", categories=[category], encoding="utf-8"
  96. )
  97. assert len(res.filenames) == 1
  98. assert len(res.target_names) == 1
  99. assert res.DESCR == "test"
  100. assert res.data == ["Hello World!\n"]
  101. def test_load_files_wo_load_content(
  102. test_category_dir_1, test_category_dir_2, load_files_root
  103. ):
  104. res = load_files(load_files_root, load_content=False)
  105. assert len(res.filenames) == 1
  106. assert len(res.target_names) == 2
  107. assert res.DESCR is None
  108. assert res.get("data") is None
  109. @pytest.mark.parametrize("allowed_extensions", ([".txt"], [".txt", ".json"]))
  110. def test_load_files_allowed_extensions(tmp_path, allowed_extensions):
  111. """Check the behaviour of `allowed_extension` in `load_files`."""
  112. d = tmp_path / "sub"
  113. d.mkdir()
  114. files = ("file1.txt", "file2.json", "file3.json", "file4.md")
  115. paths = [d / f for f in files]
  116. for p in paths:
  117. p.write_bytes(b"hello")
  118. res = load_files(tmp_path, allowed_extensions=allowed_extensions)
  119. assert set([str(p) for p in paths if p.suffix in allowed_extensions]) == set(
  120. res.filenames
  121. )
  122. @pytest.mark.parametrize(
  123. "filename, expected_n_samples, expected_n_features, expected_target_names",
  124. [
  125. ("wine_data.csv", 178, 13, ["class_0", "class_1", "class_2"]),
  126. ("iris.csv", 150, 4, ["setosa", "versicolor", "virginica"]),
  127. ("breast_cancer.csv", 569, 30, ["malignant", "benign"]),
  128. ],
  129. )
  130. def test_load_csv_data(
  131. filename, expected_n_samples, expected_n_features, expected_target_names
  132. ):
  133. actual_data, actual_target, actual_target_names = load_csv_data(filename)
  134. assert actual_data.shape[0] == expected_n_samples
  135. assert actual_data.shape[1] == expected_n_features
  136. assert actual_target.shape[0] == expected_n_samples
  137. np.testing.assert_array_equal(actual_target_names, expected_target_names)
  138. def test_load_csv_data_with_descr():
  139. data_file_name = "iris.csv"
  140. descr_file_name = "iris.rst"
  141. res_without_descr = load_csv_data(data_file_name=data_file_name)
  142. res_with_descr = load_csv_data(
  143. data_file_name=data_file_name, descr_file_name=descr_file_name
  144. )
  145. assert len(res_with_descr) == 4
  146. assert len(res_without_descr) == 3
  147. np.testing.assert_array_equal(res_with_descr[0], res_without_descr[0])
  148. np.testing.assert_array_equal(res_with_descr[1], res_without_descr[1])
  149. np.testing.assert_array_equal(res_with_descr[2], res_without_descr[2])
  150. assert res_with_descr[-1].startswith(".. _iris_dataset:")
  151. @pytest.mark.parametrize(
  152. "filename, kwargs, expected_shape",
  153. [
  154. ("diabetes_data_raw.csv.gz", {}, [442, 10]),
  155. ("diabetes_target.csv.gz", {}, [442]),
  156. ("digits.csv.gz", {"delimiter": ","}, [1797, 65]),
  157. ],
  158. )
  159. def test_load_gzip_compressed_csv_data(filename, kwargs, expected_shape):
  160. actual_data = load_gzip_compressed_csv_data(filename, **kwargs)
  161. assert actual_data.shape == tuple(expected_shape)
  162. def test_load_gzip_compressed_csv_data_with_descr():
  163. data_file_name = "diabetes_target.csv.gz"
  164. descr_file_name = "diabetes.rst"
  165. expected_data = load_gzip_compressed_csv_data(data_file_name=data_file_name)
  166. actual_data, descr = load_gzip_compressed_csv_data(
  167. data_file_name=data_file_name,
  168. descr_file_name=descr_file_name,
  169. )
  170. np.testing.assert_array_equal(actual_data, expected_data)
  171. assert descr.startswith(".. _diabetes_dataset:")
  172. def test_load_sample_images():
  173. try:
  174. res = load_sample_images()
  175. assert len(res.images) == 2
  176. assert len(res.filenames) == 2
  177. images = res.images
  178. # assert is china image
  179. assert np.all(images[0][0, 0, :] == np.array([174, 201, 231], dtype=np.uint8))
  180. # assert is flower image
  181. assert np.all(images[1][0, 0, :] == np.array([2, 19, 13], dtype=np.uint8))
  182. assert res.DESCR
  183. except ImportError:
  184. warnings.warn("Could not load sample images, PIL is not available.")
  185. def test_load_sample_image():
  186. try:
  187. china = load_sample_image("china.jpg")
  188. assert china.dtype == "uint8"
  189. assert china.shape == (427, 640, 3)
  190. except ImportError:
  191. warnings.warn("Could not load sample images, PIL is not available.")
  192. def test_load_diabetes_raw():
  193. """Test to check that we load a scaled version by default but that we can
  194. get an unscaled version when setting `scaled=False`."""
  195. diabetes_raw = load_diabetes(scaled=False)
  196. assert diabetes_raw.data.shape == (442, 10)
  197. assert diabetes_raw.target.size, 442
  198. assert len(diabetes_raw.feature_names) == 10
  199. assert diabetes_raw.DESCR
  200. diabetes_default = load_diabetes()
  201. np.testing.assert_allclose(
  202. scale(diabetes_raw.data) / (442**0.5), diabetes_default.data, atol=1e-04
  203. )
  204. @pytest.mark.parametrize(
  205. "loader_func, data_shape, target_shape, n_target, has_descr, filenames",
  206. [
  207. (load_breast_cancer, (569, 30), (569,), 2, True, ["filename"]),
  208. (load_wine, (178, 13), (178,), 3, True, []),
  209. (load_iris, (150, 4), (150,), 3, True, ["filename"]),
  210. (
  211. load_linnerud,
  212. (20, 3),
  213. (20, 3),
  214. 3,
  215. True,
  216. ["data_filename", "target_filename"],
  217. ),
  218. (load_diabetes, (442, 10), (442,), None, True, []),
  219. (load_digits, (1797, 64), (1797,), 10, True, []),
  220. (partial(load_digits, n_class=9), (1617, 64), (1617,), 10, True, []),
  221. ],
  222. )
  223. def test_loader(loader_func, data_shape, target_shape, n_target, has_descr, filenames):
  224. bunch = loader_func()
  225. assert isinstance(bunch, Bunch)
  226. assert bunch.data.shape == data_shape
  227. assert bunch.target.shape == target_shape
  228. if hasattr(bunch, "feature_names"):
  229. assert len(bunch.feature_names) == data_shape[1]
  230. if n_target is not None:
  231. assert len(bunch.target_names) == n_target
  232. if has_descr:
  233. assert bunch.DESCR
  234. if filenames:
  235. assert "data_module" in bunch
  236. assert all(
  237. [
  238. f in bunch and _is_resource(bunch["data_module"], bunch[f])
  239. for f in filenames
  240. ]
  241. )
  242. @pytest.mark.parametrize(
  243. "loader_func, data_dtype, target_dtype",
  244. [
  245. (load_breast_cancer, np.float64, int),
  246. (load_diabetes, np.float64, np.float64),
  247. (load_digits, np.float64, int),
  248. (load_iris, np.float64, int),
  249. (load_linnerud, np.float64, np.float64),
  250. (load_wine, np.float64, int),
  251. ],
  252. )
  253. def test_toy_dataset_frame_dtype(loader_func, data_dtype, target_dtype):
  254. default_result = loader_func()
  255. check_as_frame(
  256. default_result,
  257. loader_func,
  258. expected_data_dtype=data_dtype,
  259. expected_target_dtype=target_dtype,
  260. )
  261. def test_loads_dumps_bunch():
  262. bunch = Bunch(x="x")
  263. bunch_from_pkl = loads(dumps(bunch))
  264. bunch_from_pkl.x = "y"
  265. assert bunch_from_pkl["x"] == bunch_from_pkl.x
  266. def test_bunch_pickle_generated_with_0_16_and_read_with_0_17():
  267. bunch = Bunch(key="original")
  268. # This reproduces a problem when Bunch pickles have been created
  269. # with scikit-learn 0.16 and are read with 0.17. Basically there
  270. # is a surprising behaviour because reading bunch.key uses
  271. # bunch.__dict__ (which is non empty for 0.16 Bunch objects)
  272. # whereas assigning into bunch.key uses bunch.__setattr__. See
  273. # https://github.com/scikit-learn/scikit-learn/issues/6196 for
  274. # more details
  275. bunch.__dict__["key"] = "set from __dict__"
  276. bunch_from_pkl = loads(dumps(bunch))
  277. # After loading from pickle the __dict__ should have been ignored
  278. assert bunch_from_pkl.key == "original"
  279. assert bunch_from_pkl["key"] == "original"
  280. # Making sure that changing the attr does change the value
  281. # associated with __getitem__ as well
  282. bunch_from_pkl.key = "changed"
  283. assert bunch_from_pkl.key == "changed"
  284. assert bunch_from_pkl["key"] == "changed"
  285. def test_bunch_dir():
  286. # check that dir (important for autocomplete) shows attributes
  287. data = load_iris()
  288. assert "data" in dir(data)
  289. def test_load_boston_error():
  290. """Check that we raise the ethical warning when trying to import `load_boston`."""
  291. msg = "The Boston housing prices dataset has an ethical problem"
  292. with pytest.raises(ImportError, match=msg):
  293. from sklearn.datasets import load_boston # noqa
  294. # other non-existing function should raise the usual import error
  295. msg = "cannot import name 'non_existing_function' from 'sklearn.datasets'"
  296. with pytest.raises(ImportError, match=msg):
  297. from sklearn.datasets import non_existing_function # noqa