Still working on file handlers... JsonWriter (#2)
This commit is contained in:
@ -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',
|
||||||
|
|||||||
@ -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
|
||||||
|
|
||||||
|
|||||||
@ -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)
|
||||||
|
|||||||
@ -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')
|
||||||
|
|||||||
Reference in New Issue
Block a user