_datasource.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702
  1. """A file interface for handling local and remote data files.
  2. The goal of datasource is to abstract some of the file system operations
  3. when dealing with data files so the researcher doesn't have to know all the
  4. low-level details. Through datasource, a researcher can obtain and use a
  5. file with one function call, regardless of location of the file.
  6. DataSource is meant to augment standard python libraries, not replace them.
  7. It should work seamlessly with standard file IO operations and the os
  8. module.
  9. DataSource files can originate locally or remotely:
  10. - local files : '/home/guido/src/local/data.txt'
  11. - URLs (http, ftp, ...) : 'http://www.scipy.org/not/real/data.txt'
  12. DataSource files can also be compressed or uncompressed. Currently only
  13. gzip, bz2 and xz are supported.
  14. Example::
  15. >>> # Create a DataSource, use os.curdir (default) for local storage.
  16. >>> from numpy import DataSource
  17. >>> ds = DataSource()
  18. >>>
  19. >>> # Open a remote file.
  20. >>> # DataSource downloads the file, stores it locally in:
  21. >>> # './www.google.com/index.html'
  22. >>> # opens the file and returns a file object.
  23. >>> fp = ds.open('http://www.google.com/') # doctest: +SKIP
  24. >>>
  25. >>> # Use the file as you normally would
  26. >>> fp.read() # doctest: +SKIP
  27. >>> fp.close() # doctest: +SKIP
  28. """
  29. import os
  30. import shutil
  31. import io
  32. from numpy.core.overrides import set_module
  33. _open = open
  34. def _check_mode(mode, encoding, newline):
  35. """Check mode and that encoding and newline are compatible.
  36. Parameters
  37. ----------
  38. mode : str
  39. File open mode.
  40. encoding : str
  41. File encoding.
  42. newline : str
  43. Newline for text files.
  44. """
  45. if "t" in mode:
  46. if "b" in mode:
  47. raise ValueError("Invalid mode: %r" % (mode,))
  48. else:
  49. if encoding is not None:
  50. raise ValueError("Argument 'encoding' not supported in binary mode")
  51. if newline is not None:
  52. raise ValueError("Argument 'newline' not supported in binary mode")
  53. # Using a class instead of a module-level dictionary
  54. # to reduce the initial 'import numpy' overhead by
  55. # deferring the import of lzma, bz2 and gzip until needed
  56. # TODO: .zip support, .tar support?
  57. class _FileOpeners:
  58. """
  59. Container for different methods to open (un-)compressed files.
  60. `_FileOpeners` contains a dictionary that holds one method for each
  61. supported file format. Attribute lookup is implemented in such a way
  62. that an instance of `_FileOpeners` itself can be indexed with the keys
  63. of that dictionary. Currently uncompressed files as well as files
  64. compressed with ``gzip``, ``bz2`` or ``xz`` compression are supported.
  65. Notes
  66. -----
  67. `_file_openers`, an instance of `_FileOpeners`, is made available for
  68. use in the `_datasource` module.
  69. Examples
  70. --------
  71. >>> import gzip
  72. >>> np.lib._datasource._file_openers.keys()
  73. [None, '.bz2', '.gz', '.xz', '.lzma']
  74. >>> np.lib._datasource._file_openers['.gz'] is gzip.open
  75. True
  76. """
  77. def __init__(self):
  78. self._loaded = False
  79. self._file_openers = {None: io.open}
  80. def _load(self):
  81. if self._loaded:
  82. return
  83. try:
  84. import bz2
  85. self._file_openers[".bz2"] = bz2.open
  86. except ImportError:
  87. pass
  88. try:
  89. import gzip
  90. self._file_openers[".gz"] = gzip.open
  91. except ImportError:
  92. pass
  93. try:
  94. import lzma
  95. self._file_openers[".xz"] = lzma.open
  96. self._file_openers[".lzma"] = lzma.open
  97. except (ImportError, AttributeError):
  98. # There are incompatible backports of lzma that do not have the
  99. # lzma.open attribute, so catch that as well as ImportError.
  100. pass
  101. self._loaded = True
  102. def keys(self):
  103. """
  104. Return the keys of currently supported file openers.
  105. Parameters
  106. ----------
  107. None
  108. Returns
  109. -------
  110. keys : list
  111. The keys are None for uncompressed files and the file extension
  112. strings (i.e. ``'.gz'``, ``'.xz'``) for supported compression
  113. methods.
  114. """
  115. self._load()
  116. return list(self._file_openers.keys())
  117. def __getitem__(self, key):
  118. self._load()
  119. return self._file_openers[key]
  120. _file_openers = _FileOpeners()
  121. def open(path, mode='r', destpath=os.curdir, encoding=None, newline=None):
  122. """
  123. Open `path` with `mode` and return the file object.
  124. If ``path`` is an URL, it will be downloaded, stored in the
  125. `DataSource` `destpath` directory and opened from there.
  126. Parameters
  127. ----------
  128. path : str
  129. Local file path or URL to open.
  130. mode : str, optional
  131. Mode to open `path`. Mode 'r' for reading, 'w' for writing, 'a' to
  132. append. Available modes depend on the type of object specified by
  133. path. Default is 'r'.
  134. destpath : str, optional
  135. Path to the directory where the source file gets downloaded to for
  136. use. If `destpath` is None, a temporary directory will be created.
  137. The default path is the current directory.
  138. encoding : {None, str}, optional
  139. Open text file with given encoding. The default encoding will be
  140. what `io.open` uses.
  141. newline : {None, str}, optional
  142. Newline to use when reading text file.
  143. Returns
  144. -------
  145. out : file object
  146. The opened file.
  147. Notes
  148. -----
  149. This is a convenience function that instantiates a `DataSource` and
  150. returns the file object from ``DataSource.open(path)``.
  151. """
  152. ds = DataSource(destpath)
  153. return ds.open(path, mode, encoding=encoding, newline=newline)
  154. @set_module('numpy')
  155. class DataSource:
  156. """
  157. DataSource(destpath='.')
  158. A generic data source file (file, http, ftp, ...).
  159. DataSources can be local files or remote files/URLs. The files may
  160. also be compressed or uncompressed. DataSource hides some of the
  161. low-level details of downloading the file, allowing you to simply pass
  162. in a valid file path (or URL) and obtain a file object.
  163. Parameters
  164. ----------
  165. destpath : str or None, optional
  166. Path to the directory where the source file gets downloaded to for
  167. use. If `destpath` is None, a temporary directory will be created.
  168. The default path is the current directory.
  169. Notes
  170. -----
  171. URLs require a scheme string (``http://``) to be used, without it they
  172. will fail::
  173. >>> repos = np.DataSource()
  174. >>> repos.exists('www.google.com/index.html')
  175. False
  176. >>> repos.exists('http://www.google.com/index.html')
  177. True
  178. Temporary directories are deleted when the DataSource is deleted.
  179. Examples
  180. --------
  181. ::
  182. >>> ds = np.DataSource('/home/guido')
  183. >>> urlname = 'http://www.google.com/'
  184. >>> gfile = ds.open('http://www.google.com/')
  185. >>> ds.abspath(urlname)
  186. '/home/guido/www.google.com/index.html'
  187. >>> ds = np.DataSource(None) # use with temporary file
  188. >>> ds.open('/home/guido/foobar.txt')
  189. <open file '/home/guido.foobar.txt', mode 'r' at 0x91d4430>
  190. >>> ds.abspath('/home/guido/foobar.txt')
  191. '/tmp/.../home/guido/foobar.txt'
  192. """
  193. def __init__(self, destpath=os.curdir):
  194. """Create a DataSource with a local path at destpath."""
  195. if destpath:
  196. self._destpath = os.path.abspath(destpath)
  197. self._istmpdest = False
  198. else:
  199. import tempfile # deferring import to improve startup time
  200. self._destpath = tempfile.mkdtemp()
  201. self._istmpdest = True
  202. def __del__(self):
  203. # Remove temp directories
  204. if hasattr(self, '_istmpdest') and self._istmpdest:
  205. shutil.rmtree(self._destpath)
  206. def _iszip(self, filename):
  207. """Test if the filename is a zip file by looking at the file extension.
  208. """
  209. fname, ext = os.path.splitext(filename)
  210. return ext in _file_openers.keys()
  211. def _iswritemode(self, mode):
  212. """Test if the given mode will open a file for writing."""
  213. # Currently only used to test the bz2 files.
  214. _writemodes = ("w", "+")
  215. for c in mode:
  216. if c in _writemodes:
  217. return True
  218. return False
  219. def _splitzipext(self, filename):
  220. """Split zip extension from filename and return filename.
  221. *Returns*:
  222. base, zip_ext : {tuple}
  223. """
  224. if self._iszip(filename):
  225. return os.path.splitext(filename)
  226. else:
  227. return filename, None
  228. def _possible_names(self, filename):
  229. """Return a tuple containing compressed filename variations."""
  230. names = [filename]
  231. if not self._iszip(filename):
  232. for zipext in _file_openers.keys():
  233. if zipext:
  234. names.append(filename+zipext)
  235. return names
  236. def _isurl(self, path):
  237. """Test if path is a net location. Tests the scheme and netloc."""
  238. # We do this here to reduce the 'import numpy' initial import time.
  239. from urllib.parse import urlparse
  240. # BUG : URLs require a scheme string ('http://') to be used.
  241. # www.google.com will fail.
  242. # Should we prepend the scheme for those that don't have it and
  243. # test that also? Similar to the way we append .gz and test for
  244. # for compressed versions of files.
  245. scheme, netloc, upath, uparams, uquery, ufrag = urlparse(path)
  246. return bool(scheme and netloc)
  247. def _cache(self, path):
  248. """Cache the file specified by path.
  249. Creates a copy of the file in the datasource cache.
  250. """
  251. # We import these here because importing urllib is slow and
  252. # a significant fraction of numpy's total import time.
  253. from urllib.request import urlopen
  254. from urllib.error import URLError
  255. upath = self.abspath(path)
  256. # ensure directory exists
  257. if not os.path.exists(os.path.dirname(upath)):
  258. os.makedirs(os.path.dirname(upath))
  259. # TODO: Doesn't handle compressed files!
  260. if self._isurl(path):
  261. with urlopen(path) as openedurl:
  262. with _open(upath, 'wb') as f:
  263. shutil.copyfileobj(openedurl, f)
  264. else:
  265. shutil.copyfile(path, upath)
  266. return upath
  267. def _findfile(self, path):
  268. """Searches for ``path`` and returns full path if found.
  269. If path is an URL, _findfile will cache a local copy and return the
  270. path to the cached file. If path is a local file, _findfile will
  271. return a path to that local file.
  272. The search will include possible compressed versions of the file
  273. and return the first occurrence found.
  274. """
  275. # Build list of possible local file paths
  276. if not self._isurl(path):
  277. # Valid local paths
  278. filelist = self._possible_names(path)
  279. # Paths in self._destpath
  280. filelist += self._possible_names(self.abspath(path))
  281. else:
  282. # Cached URLs in self._destpath
  283. filelist = self._possible_names(self.abspath(path))
  284. # Remote URLs
  285. filelist = filelist + self._possible_names(path)
  286. for name in filelist:
  287. if self.exists(name):
  288. if self._isurl(name):
  289. name = self._cache(name)
  290. return name
  291. return None
  292. def abspath(self, path):
  293. """
  294. Return absolute path of file in the DataSource directory.
  295. If `path` is an URL, then `abspath` will return either the location
  296. the file exists locally or the location it would exist when opened
  297. using the `open` method.
  298. Parameters
  299. ----------
  300. path : str
  301. Can be a local file or a remote URL.
  302. Returns
  303. -------
  304. out : str
  305. Complete path, including the `DataSource` destination directory.
  306. Notes
  307. -----
  308. The functionality is based on `os.path.abspath`.
  309. """
  310. # We do this here to reduce the 'import numpy' initial import time.
  311. from urllib.parse import urlparse
  312. # TODO: This should be more robust. Handles case where path includes
  313. # the destpath, but not other sub-paths. Failing case:
  314. # path = /home/guido/datafile.txt
  315. # destpath = /home/alex/
  316. # upath = self.abspath(path)
  317. # upath == '/home/alex/home/guido/datafile.txt'
  318. # handle case where path includes self._destpath
  319. splitpath = path.split(self._destpath, 2)
  320. if len(splitpath) > 1:
  321. path = splitpath[1]
  322. scheme, netloc, upath, uparams, uquery, ufrag = urlparse(path)
  323. netloc = self._sanitize_relative_path(netloc)
  324. upath = self._sanitize_relative_path(upath)
  325. return os.path.join(self._destpath, netloc, upath)
  326. def _sanitize_relative_path(self, path):
  327. """Return a sanitised relative path for which
  328. os.path.abspath(os.path.join(base, path)).startswith(base)
  329. """
  330. last = None
  331. path = os.path.normpath(path)
  332. while path != last:
  333. last = path
  334. # Note: os.path.join treats '/' as os.sep on Windows
  335. path = path.lstrip(os.sep).lstrip('/')
  336. path = path.lstrip(os.pardir).lstrip('..')
  337. drive, path = os.path.splitdrive(path) # for Windows
  338. return path
  339. def exists(self, path):
  340. """
  341. Test if path exists.
  342. Test if `path` exists as (and in this order):
  343. - a local file.
  344. - a remote URL that has been downloaded and stored locally in the
  345. `DataSource` directory.
  346. - a remote URL that has not been downloaded, but is valid and
  347. accessible.
  348. Parameters
  349. ----------
  350. path : str
  351. Can be a local file or a remote URL.
  352. Returns
  353. -------
  354. out : bool
  355. True if `path` exists.
  356. Notes
  357. -----
  358. When `path` is an URL, `exists` will return True if it's either
  359. stored locally in the `DataSource` directory, or is a valid remote
  360. URL. `DataSource` does not discriminate between the two, the file
  361. is accessible if it exists in either location.
  362. """
  363. # First test for local path
  364. if os.path.exists(path):
  365. return True
  366. # We import this here because importing urllib is slow and
  367. # a significant fraction of numpy's total import time.
  368. from urllib.request import urlopen
  369. from urllib.error import URLError
  370. # Test cached url
  371. upath = self.abspath(path)
  372. if os.path.exists(upath):
  373. return True
  374. # Test remote url
  375. if self._isurl(path):
  376. try:
  377. netfile = urlopen(path)
  378. netfile.close()
  379. del(netfile)
  380. return True
  381. except URLError:
  382. return False
  383. return False
  384. def open(self, path, mode='r', encoding=None, newline=None):
  385. """
  386. Open and return file-like object.
  387. If `path` is an URL, it will be downloaded, stored in the
  388. `DataSource` directory and opened from there.
  389. Parameters
  390. ----------
  391. path : str
  392. Local file path or URL to open.
  393. mode : {'r', 'w', 'a'}, optional
  394. Mode to open `path`. Mode 'r' for reading, 'w' for writing,
  395. 'a' to append. Available modes depend on the type of object
  396. specified by `path`. Default is 'r'.
  397. encoding : {None, str}, optional
  398. Open text file with given encoding. The default encoding will be
  399. what `io.open` uses.
  400. newline : {None, str}, optional
  401. Newline to use when reading text file.
  402. Returns
  403. -------
  404. out : file object
  405. File object.
  406. """
  407. # TODO: There is no support for opening a file for writing which
  408. # doesn't exist yet (creating a file). Should there be?
  409. # TODO: Add a ``subdir`` parameter for specifying the subdirectory
  410. # used to store URLs in self._destpath.
  411. if self._isurl(path) and self._iswritemode(mode):
  412. raise ValueError("URLs are not writeable")
  413. # NOTE: _findfile will fail on a new file opened for writing.
  414. found = self._findfile(path)
  415. if found:
  416. _fname, ext = self._splitzipext(found)
  417. if ext == 'bz2':
  418. mode.replace("+", "")
  419. return _file_openers[ext](found, mode=mode,
  420. encoding=encoding, newline=newline)
  421. else:
  422. raise IOError("%s not found." % path)
  423. class Repository (DataSource):
  424. """
  425. Repository(baseurl, destpath='.')
  426. A data repository where multiple DataSource's share a base
  427. URL/directory.
  428. `Repository` extends `DataSource` by prepending a base URL (or
  429. directory) to all the files it handles. Use `Repository` when you will
  430. be working with multiple files from one base URL. Initialize
  431. `Repository` with the base URL, then refer to each file by its filename
  432. only.
  433. Parameters
  434. ----------
  435. baseurl : str
  436. Path to the local directory or remote location that contains the
  437. data files.
  438. destpath : str or None, optional
  439. Path to the directory where the source file gets downloaded to for
  440. use. If `destpath` is None, a temporary directory will be created.
  441. The default path is the current directory.
  442. Examples
  443. --------
  444. To analyze all files in the repository, do something like this
  445. (note: this is not self-contained code)::
  446. >>> repos = np.lib._datasource.Repository('/home/user/data/dir/')
  447. >>> for filename in filelist:
  448. ... fp = repos.open(filename)
  449. ... fp.analyze()
  450. ... fp.close()
  451. Similarly you could use a URL for a repository::
  452. >>> repos = np.lib._datasource.Repository('http://www.xyz.edu/data')
  453. """
  454. def __init__(self, baseurl, destpath=os.curdir):
  455. """Create a Repository with a shared url or directory of baseurl."""
  456. DataSource.__init__(self, destpath=destpath)
  457. self._baseurl = baseurl
  458. def __del__(self):
  459. DataSource.__del__(self)
  460. def _fullpath(self, path):
  461. """Return complete path for path. Prepends baseurl if necessary."""
  462. splitpath = path.split(self._baseurl, 2)
  463. if len(splitpath) == 1:
  464. result = os.path.join(self._baseurl, path)
  465. else:
  466. result = path # path contains baseurl already
  467. return result
  468. def _findfile(self, path):
  469. """Extend DataSource method to prepend baseurl to ``path``."""
  470. return DataSource._findfile(self, self._fullpath(path))
  471. def abspath(self, path):
  472. """
  473. Return absolute path of file in the Repository directory.
  474. If `path` is an URL, then `abspath` will return either the location
  475. the file exists locally or the location it would exist when opened
  476. using the `open` method.
  477. Parameters
  478. ----------
  479. path : str
  480. Can be a local file or a remote URL. This may, but does not
  481. have to, include the `baseurl` with which the `Repository` was
  482. initialized.
  483. Returns
  484. -------
  485. out : str
  486. Complete path, including the `DataSource` destination directory.
  487. """
  488. return DataSource.abspath(self, self._fullpath(path))
  489. def exists(self, path):
  490. """
  491. Test if path exists prepending Repository base URL to path.
  492. Test if `path` exists as (and in this order):
  493. - a local file.
  494. - a remote URL that has been downloaded and stored locally in the
  495. `DataSource` directory.
  496. - a remote URL that has not been downloaded, but is valid and
  497. accessible.
  498. Parameters
  499. ----------
  500. path : str
  501. Can be a local file or a remote URL. This may, but does not
  502. have to, include the `baseurl` with which the `Repository` was
  503. initialized.
  504. Returns
  505. -------
  506. out : bool
  507. True if `path` exists.
  508. Notes
  509. -----
  510. When `path` is an URL, `exists` will return True if it's either
  511. stored locally in the `DataSource` directory, or is a valid remote
  512. URL. `DataSource` does not discriminate between the two, the file
  513. is accessible if it exists in either location.
  514. """
  515. return DataSource.exists(self, self._fullpath(path))
  516. def open(self, path, mode='r', encoding=None, newline=None):
  517. """
  518. Open and return file-like object prepending Repository base URL.
  519. If `path` is an URL, it will be downloaded, stored in the
  520. DataSource directory and opened from there.
  521. Parameters
  522. ----------
  523. path : str
  524. Local file path or URL to open. This may, but does not have to,
  525. include the `baseurl` with which the `Repository` was
  526. initialized.
  527. mode : {'r', 'w', 'a'}, optional
  528. Mode to open `path`. Mode 'r' for reading, 'w' for writing,
  529. 'a' to append. Available modes depend on the type of object
  530. specified by `path`. Default is 'r'.
  531. encoding : {None, str}, optional
  532. Open text file with given encoding. The default encoding will be
  533. what `io.open` uses.
  534. newline : {None, str}, optional
  535. Newline to use when reading text file.
  536. Returns
  537. -------
  538. out : file object
  539. File object.
  540. """
  541. return DataSource.open(self, self._fullpath(path), mode,
  542. encoding=encoding, newline=newline)
  543. def listdir(self):
  544. """
  545. List files in the source Repository.
  546. Returns
  547. -------
  548. files : list of str
  549. List of file names (not containing a directory part).
  550. Notes
  551. -----
  552. Does not currently work for remote repositories.
  553. """
  554. if self._isurl(self._baseurl):
  555. raise NotImplementedError(
  556. "Directory listing of URLs, not supported yet.")
  557. else:
  558. return os.listdir(self._baseurl)