diff --git a/pyinfra/visitor/strategies/response/aggregation.py b/pyinfra/visitor/strategies/response/aggregation.py index 928b389..c0f96af 100644 --- a/pyinfra/visitor/strategies/response/aggregation.py +++ b/pyinfra/visitor/strategies/response/aggregation.py @@ -1,7 +1,7 @@ from collections import deque from typing import Callable -from funcy import omit, filter +from funcy import omit, filter, first from more_itertools import peekable from pyinfra.server.dispatcher.dispatcher import Nothing, is_not_nothing @@ -10,10 +10,16 @@ from pyinfra.visitor.strategies.response.response import ResponseStrategy from pyinfra.visitor.utils import build_storage_upload_info +def default_merge(items): + merged = list(items) + merged = first(merged) if len(merged) == 1 else merged + return merged + + class AggregationStorageStrategy(ResponseStrategy): def __init__(self, storage, merger: Callable = None): self.storage = storage - self.merger = merger or list + self.merger = merger or default_merge self.buffer = deque() self.response_files = deque() diff --git a/test/integration_tests/serve_test.py b/test/integration_tests/serve_test.py index aa56497..2eeffd3 100644 --- a/test/integration_tests/serve_test.py +++ b/test/integration_tests/serve_test.py @@ -4,7 +4,7 @@ from itertools import starmap, repeat, chain from operator import itemgetter import pytest -from funcy import compose, lpluck, first, second, pluck, lflatten +from funcy import compose, first, second, pluck, lflatten from pyinfra.default_objects import ( get_callback,