OpenOrStdin/out: add encoding

This commit is contained in:
Doug Blank 2015-08-11 15:03:27 -04:00
parent dbb1644136
commit e5a77a60f6

View File

@ -379,13 +379,16 @@ def load_addon_file(path, callback=None):
#------------------------------------------------------------------------- #-------------------------------------------------------------------------
class OpenFileOrStdout: class OpenFileOrStdout:
"""Context manager to open a file or stdout for writing.""" """Context manager to open a file or stdout for writing."""
def __init__(self, filename): def __init__(self, filename, encoding=None):
self.filename = filename self.filename = filename
self.filehandle = None self.filehandle = None
self.encoding = encoding
def __enter__(self): def __enter__(self):
if self.filename == '-': if self.filename == '-':
self.filehandle = sys.stdout self.filehandle = sys.stdout
elif self.encoding:
self.filehandle = open(self.filename, 'w', encoding=self.encoding)
else: else:
self.filehandle = open(self.filename, 'w') self.filehandle = open(self.filename, 'w')
return self.filehandle return self.filehandle
@ -402,14 +405,17 @@ class OpenFileOrStdout:
#------------------------------------------------------------------------- #-------------------------------------------------------------------------
class OpenFileOrStdin: class OpenFileOrStdin:
"""Context manager to open a file or stdin for reading.""" """Context manager to open a file or stdin for reading."""
def __init__(self, filename, add_mode=''): def __init__(self, filename, add_mode='', encoding=None):
self.filename = filename self.filename = filename
self.mode = 'r%s' % add_mode self.mode = 'r%s' % add_mode
self.filehandle = None self.filehandle = None
self.encoding = encoding
def __enter__(self): def __enter__(self):
if self.filename == '-': if self.filename == '-':
self.filehandle = sys.stdin self.filehandle = sys.stdin
elif self.encoding:
self.filehandle = open(self.filename, self.mode, encoding=self.encoding)
else: else:
self.filehandle = open(self.filename, self.mode) self.filehandle = open(self.filename, self.mode)
return self.filehandle return self.filehandle