Created
April 14, 2014 16:14
-
-
Save Jorge-C/10662095 to your computer and use it in GitHub Desktop.
open++
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import unittest | |
| import tempfile | |
| from StringIO import StringIO | |
| from util import open_filepath_or, _is_string_or_bytes | |
| class TestFilePathOpening(unittest.TestCase): | |
| def test_is_string_or_bytes(self): | |
| self.assertTrue(_is_string_or_bytes('foo')) | |
| self.assertTrue(_is_string_or_bytes(u'foo')) | |
| self.assertTrue(_is_string_or_bytes(b'foo')) | |
| self.assertFalse(_is_string_or_bytes(StringIO('bar'))) | |
| self.assertFalse(_is_string_or_bytes([1])) | |
| def test_file_closed(self): | |
| """File gets closed in decorator""" | |
| f = tempfile.NamedTemporaryFile('r') | |
| filepath = f.name | |
| with open_filepath_or(filepath) as fh: | |
| pass | |
| self.assertTrue(fh.closed) | |
| def test_file_closed_harder(self): | |
| """File gets closed in decorator, even if exceptions happen.""" | |
| f = tempfile.NamedTemporaryFile('r') | |
| filepath = f.name | |
| try: | |
| with open_filepath_or(filepath) as fh: | |
| raise TypeError | |
| except TypeError: | |
| self.assertTrue(fh.closed) | |
| else: | |
| # If we're here, no exceptions have been raised inside the | |
| # try clause, so the context manager swallowed them. No | |
| # good. | |
| raise Exception("`open_filepath_or` didn't propagate exceptions") | |
| def test_filehandle(self): | |
| """Filehandles slip through untouched""" | |
| with tempfile.TemporaryFile('r') as fh: | |
| with open_filepath_or(fh) as ffh: | |
| self.assertTrue(fh is ffh) | |
| # And it doesn't close the file-handle | |
| self.assertFalse(fh.closed) | |
| def test_StringIO(self): | |
| """StringIO (useful e.g. for testing) slips through.""" | |
| f = StringIO("File contents") | |
| with open_filepath_or(f) as fh: | |
| self.assertTrue(fh is f) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| from contextlib import contextmanager | |
| from future.builtins import bytes, str | |
| def _is_string_or_bytes(s): | |
| """Returns True if input argument is string (unicode or not) or bytes. | |
| """ | |
| return isinstance(s, str) or isinstance(s, bytes) | |
| def _get_filehandle(filepath_or, *args, **kwargs): | |
| """Open file if `filepath_or` looks like a string/unicode/bytes, else | |
| pass through. | |
| """ | |
| if _is_string_or_bytes(filepath_or): | |
| fh, own_fh = open(filepath_or, *args, **kwargs), True | |
| else: | |
| fh, own_fh = filepath_or, False | |
| return fh, own_fh | |
| @contextmanager | |
| def open_filepath_or(filepath_or, *args, **kwargs): | |
| """Context manager, like ``open``, but let's file handles and file | |
| like objects pass untouched. | |
| It is useful when implementing a function that can accept both | |
| strings and file-like objects (like numpy.loadtxt, etc). | |
| Parameters | |
| ---------- | |
| filepath_or : str/bytes/unicode string or file-like | |
| If string, file to be opened using ``open``. Else, it is returned | |
| untouched. | |
| Other parameters | |
| ---------------- | |
| When `filepath_or` is a string, any extra arguments are passed on to | |
| ``open``. | |
| Examples | |
| -------- | |
| >>> with open_filepath_or('filename') as f: # doctest: +SKIP | |
| ... pass | |
| >>> fh = open('filename') # doctest: +SKIP | |
| >>> with open_filepath_or(fh) as f: # doctest: +SKIP | |
| ... pass | |
| >>> fh.closed # doctest: +SKIP | |
| False | |
| >>> fh.close() # doctest: +SKIP | |
| """ | |
| fh, own_fh = _get_filehandle(filepath_or, *args, **kwargs) | |
| try: | |
| yield fh | |
| finally: | |
| if own_fh: | |
| fh.close() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment