Skip to content

Commit

Permalink
First stab at generic exception handling
Browse files Browse the repository at this point in the history
  • Loading branch information
freol35241 committed Oct 1, 2021
1 parent 14ee10f commit 4e53fac
Showing 1 changed file with 41 additions and 19 deletions.
60 changes: 41 additions & 19 deletions streamz/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,9 @@ def __init__(self, upstream=None, upstreams=None, stream_name=None,
else:
self.upstreams = []

# Lazily loaded exception handler to avoid recursion
self._on_exception = None

self._set_asynchronous(asynchronous)
self._set_loop(loop)
if ensure_io_loop and not self.loop:
Expand Down Expand Up @@ -445,14 +448,18 @@ def _emit(self, x, metadata=None):

result = []
for downstream in list(self.downstreams):
r = downstream.update(x, who=self, metadata=metadata)

if type(r) is list:
result.extend(r)
else:
result.append(r)
try:
r = downstream.update(x, who=self, metadata=metadata)

self._release_refs(metadata)
if type(r) is list:
result.extend(r)
else:
result.append(r)
except Exception as exc:
# Push this exception to the on_exception handler on the downstream that raised
downstream.on_exception().update((x, exc) , who=self, metadata=metadata)
finally:
self._release_refs(metadata)

return [element for element in result if element is not None]

Expand Down Expand Up @@ -671,6 +678,30 @@ def _release_refs(self, metadata, n=1):
if 'ref' in m:
m['ref'].release(n)

def on_exception(self):
"""Returns the exception handler associated with this stream
"""
self._on_exception = self._on_exception or _on_exception()
return self._on_exception


class InvalidDataError(Exception):
pass

class _on_exception(Stream):

def __init__(self, *args, **kwargs):
self.silent = False
Stream.__init__(self, *args, **kwargs)

def update(self, x, who=None, metadata=None):
cause, exc = x

if self.silent or len(self.downstreams) > 0:
self._emit(x, metadata=metadata)
else:
logger.exception(exc)
raise InvalidDataError(cause) from exc

@Stream.register_api()
class map(Stream):
Expand Down Expand Up @@ -706,13 +737,8 @@ def __init__(self, upstream, func, *args, **kwargs):
Stream.__init__(self, upstream, stream_name=stream_name)

def update(self, x, who=None, metadata=None):
try:
result = self.func(x, *self.args, **self.kwargs)
except Exception as e:
logger.exception(e)
raise
else:
return self._emit(result, metadata=metadata)
result = self.func(x, *self.args, **self.kwargs)
self._emit(result, metadata=metadata)


@Stream.register_api()
Expand Down Expand Up @@ -890,11 +916,7 @@ def update(self, x, who=None, metadata=None):
else:
return self._emit(x, metadata=metadata)
else:
try:
result = self.func(self.state, x, **self.kwargs)
except Exception as e:
logger.exception(e)
raise
result = self.func(self.state, x, **self.kwargs)
if self.returns_state:
state, result = result
else:
Expand Down

0 comments on commit 4e53fac

Please sign in to comment.