Still working on file handlers... JsonWriter (#2)

This commit is contained in:
Romain Dorgueil
2016-12-28 11:14:37 +01:00
parent 2b8216ec0d
commit 2527a1b85b
4 changed files with 43 additions and 38 deletions

View File

@ -1,10 +1,10 @@
""" Readers and writers for common file formats. """ """ Readers and writers for common file formats. """
from .file import FileHandler, FileReader, FileWriter from .file import Handler, FileReader, FileWriter
from .json import JsonWriter from .json import JsonWriter
__all__ = [ __all__ = [
'FileHandler', 'Handler',
'FileReader', 'FileReader',
'FileWriter', 'FileWriter',
'JsonWriter', 'JsonWriter',

View File

@ -1,26 +1,25 @@
from functools import partial
from bonobo.util.lifecycle import with_context from bonobo.util.lifecycle import with_context
__all__ = [ __all__ = [
'FileHandler', 'Handler',
'FileReader', 'FileReader',
'FileWriter', 'FileWriter',
] ]
@with_context @with_context
class FileHandler: class Handler:
""" """
Abstract component factory for file-related components. Abstract component factory for file-related components.
""" """
eol = '\n'
mode = None mode = None
def __init__(self, path_or_buf, eol='\n'): def __init__(self, path_or_buf, eol=None):
self.path_or_buf = path_or_buf self.path_or_buf = path_or_buf
self.eol = eol self.eol = eol or self.eol
def open(self): def open(self):
return open(self.path_or_buf, self.mode) return open(self.path_or_buf, self.mode)
@ -40,26 +39,28 @@ class FileHandler:
assert not hasattr(ctx, 'file'), 'A file pointer is already in the context... I do not know what to say...' assert not hasattr(ctx, 'file'), 'A file pointer is already in the context... I do not know what to say...'
ctx.file = self.open() ctx.file = self.open()
def __call__(self, ctx, *args):
"""
:param ComponentExecutionContext ctx:
:param mixed row:
"""
result = self.handle(ctx, *args)
try:
yield from result
except TypeError:
return
def handle(self, ctx, *args):
raise NotImplementedError('Abstract.')
def finalize(self, ctx): def finalize(self, ctx):
self.close(ctx.file) self.close(ctx.file)
del ctx.file del ctx.file
class FileReader(FileHandler): class Reader(Handler):
def __call__(self, ctx):
yield from self.handle(ctx)
def handle(self, ctx):
raise NotImplementedError('Abstract.')
class Writer(Handler):
def __call__(self, ctx, row):
return self.handle(ctx, row)
def handle(self, ctx, row):
raise NotImplementedError('Abstract.')
class FileReader(Reader):
""" """
Component factory for file-like readers. Component factory for file-like readers.
@ -70,7 +71,7 @@ class FileReader(FileHandler):
mode = 'r' mode = 'r'
def handle(self, ctx, *args): def handle(self, ctx):
""" """
Write a row on the next line of file pointed by `ctx.file`. Write a row on the next line of file pointed by `ctx.file`.
Prefix is used for newlines. Prefix is used for newlines.
@ -78,12 +79,11 @@ class FileReader(FileHandler):
:param ctx: :param ctx:
:param row: :param row:
""" """
assert not len(args)
for line in ctx.file: for line in ctx.file:
yield line.rstrip(self.eol) yield line.rstrip(self.eol)
class FileWriter(FileHandler): class FileWriter(Writer):
""" """
Component factory for file or file-like writers. Component factory for file or file-like writers.
@ -100,13 +100,12 @@ class FileWriter(FileHandler):
def handle(self, ctx, row): def handle(self, ctx, row):
""" """
Write a row on the next line of file pointed by fp. Prefix is used for newlines. Write a row on the next line of opened file in context.
:param file fp: :param file fp:
:param str row: :param str row:
:param str prefix: :param str prefix:
""" """
self.write(ctx.file, (self.eol if ctx.line else '') + row) self.write(ctx.file, (self.eol if ctx.line else '') + row)
ctx.line += 1 ctx.line += 1

View File

@ -1,23 +1,29 @@
import json import json
from .file import FileWriter from .file import FileWriter
from bonobo.util.lifecycle import with_context
__all__ = ['JsonWriter', ] __all__ = ['JsonWriter', ]
@with_context class JsonHandler:
class JsonWriter(FileWriter): eol = ',\n'
def __init__(self, path_or_buf):
super().__init__(path_or_buf, eol=',\n')
class JsonWriter(JsonHandler, FileWriter):
def initialize(self, ctx): def initialize(self, ctx):
print('EOL', self.eol)
super().initialize(ctx) super().initialize(ctx)
ctx.fp.write('[\n') ctx.file.write('[\n')
def write(self, fp, line, prefix=''): def handle(self, ctx, row):
fp.write(prefix + json.dumps(line)) """
Write a json row on the next line of file pointed by ctx.file.
:param ctx:
:param row:
"""
return super().handle(ctx, json.dumps(row))
def finalize(self, ctx): def finalize(self, ctx):
ctx.fp.write('\n]') ctx.file.write('\n]')
super().finalize(ctx) super().finalize(ctx)

View File

@ -20,7 +20,7 @@ def test_write_json_to_file(tmpdir):
]''' ]'''
with pytest.raises(AttributeError): with pytest.raises(AttributeError):
getattr(context, 'fp') getattr(context, 'file')
with pytest.raises(AttributeError): with pytest.raises(AttributeError):
getattr(context, 'first') getattr(context, 'first')