Some optimization for #97

nanohttp/decorators.py

@@ -63,39 +63,41 @@ def decorator(func):
     return decorator
 
 
-def chunked_encoding(*args, **kwargs):
-
+def chunked_encoding(trailer_field=None, trailer_value=None):
+    """Enables chunked encoding on an action
+    """
     def decorator(func):
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
+            nonlocal trailer_field, trailer_value
             context.response_headers.add_header('Transfer-Encoding', 'chunked')
-            if trailer and isinstance(trailer, dict):
-                context.response_headers.add_header(
-                    'Trailer', list(trailer.keys())[0]
-                )
+            if trailer_field:
+                if callable(trailer_field):
+                    trailer_field = trailer_field()
+                context.response_headers.add_header('Trailer', trailer_field)
             result = func(*args, **kwargs)
-
             try:
                 while True:
                     chunk = next(result)
                     yield f'{len(chunk)}\r\n{chunk}\r\n'
 
             except StopIteration:
                 yield '0\r\n'
-                if trailer and isinstance(trailer, dict):
-                    yield f'{list(trailer.keys())[0]}: '\
-                        f'{list(trailer.values())[0]}\r\n'
+                if trailer_field and trailer_value:
+                    yield f'{trailer_field}: {trailer_value}\r\n'
                 yield '\r\n'
 
             except Exception as ex:
                 yield str(ex)
                 yield '0\r\n'
                 yield '\r\n'
+
         return wrapper
 
-    trailer = args[0] if args and not callable(args[0]) else None
-    if args and callable(args[0]):
-        return decorator(args[0])
+    if callable(trailer_field):
+        func = trailer_field
+        trailer_field = None
+        return decorator(func)
 
     return decorator
 

nanohttp/tests/test_chunked_encodeing.py

@@ -16,13 +16,13 @@ def chunked():
         yield 'third chunk'
 
     @staticmethod
-    @chunked_encoding({'field': 'value'})
+    @chunked_encoding('field', 'value')
     def trailer():
         yield 'first chunk'
         yield 'second chunk'
 
     @staticmethod
-    @chunked_encoding({'field': 'value'})
+    @chunked_encoding('field', 'value')
     def error():
         yield 'first chunk'
         yield 1/0
@@ -49,7 +49,6 @@ def test_trailer(self):
             self.assertEqual(next(output_generator), '0\r\n')
             self.assertEqual(next(output_generator), 'field: value\r\n')
             self.assertEqual(next(output_generator), '\r\n')
-
             self.assertEqual(context.response_headers['trailer'], 'field')
 
     def test_error(self):