From ea9b405d2a504e30f3197947a26d15b224537630 Mon Sep 17 00:00:00 2001
From: Matthias Bisping <matthias.bisping@iqser.com>
Date: Mon, 2 May 2022 15:50:14 +0200
Subject: [PATCH] signature harminization for 1 -> 1 and 1 -> n completed

---
 pyinfra/exceptions.py                         |  4 ++
 pyinfra/rest.py                               | 49 ++++++++++++-------
 .../partial_response_test.py                  | 28 +++--------
 .../exploration_tests/pickup_endpoint_test.py |  2 -
 test/fixtures/server.py                       |  4 +-
 5 files changed, 43 insertions(+), 44 deletions(-)

diff --git a/pyinfra/exceptions.py b/pyinfra/exceptions.py
index 56405d4..ecaa06a 100644
--- a/pyinfra/exceptions.py
+++ b/pyinfra/exceptions.py
@@ -32,3 +32,7 @@ class NoSuchContainer(KeyError):
 
 class IntentionalTestException(RuntimeError):
     pass
+
+
+class UnexpectedItemType(ValueError):
+    pass
diff --git a/pyinfra/rest.py b/pyinfra/rest.py
index 698dfbc..cb58d1c 100644
--- a/pyinfra/rest.py
+++ b/pyinfra/rest.py
@@ -1,8 +1,10 @@
 from operator import itemgetter
-from typing import Iterable
+from operator import itemgetter
+from typing import Iterable, Dict, List, Callable, Union, Tuple
 
-from funcy import compose, first, flatten
+from funcy import compose, first
 
+from pyinfra.exceptions import UnexpectedItemType
 from pyinfra.utils.func import star, lift, lstarlift
 from test.utils.server import bytes_to_string, string_to_bytes
 
@@ -22,31 +24,37 @@ def bundle(data: bytes, metadata: dict):
     return package
 
 
-def normalize_up(x):
-    return [x] if isinstance(x, tuple) else x
-
-
-def normalize_down(itr):
-
-    head = first(itr)
-
-    if not head:
-        return []
-    elif isinstance(head, tuple):
-        return head, *itr
+def normalize_up(itm: Union[Tuple, Iterable]) -> Iterable:
+    if isinstance(itm, tuple):
+        return [itm]
+    elif isinstance(itm, Iterable):
+        return itm
     else:
-        return flatten((head, *itr))
+        raise UnexpectedItemType("Encountered an item that could not be normalized to a list.")
 
 
-def unpack_op_pack(operation):
+def normalize_down(itr: Iterable[Union[Tuple, Iterable]]) -> Iterable[Tuple]:
+    def normalize_item(head):
+
+        if isinstance(head, tuple):
+            return head
+        elif isinstance(head, Iterable):
+            return first(head)
+        else:
+            raise UnexpectedItemType("Encountered an item that could not be normalized to a tuple.")
+
+    return map(normalize_item, itr)
+
+
+def unpack_op_pack(operation) -> Callable[[Dict], List[Dict]]:
     return compose(lstarlift(pack), normalize_up, star(operation), unpack)
 
 
-def unpack_batchop_pack(operation):
+def unpack_batchop_pack(operation) -> Callable[[List[Dict]], List[Dict]]:
     return compose(lstarlift(pack), normalize_down, operation, lift(unpack))
 
 
-def inspect(msg="ins"):
+def inspect(msg="ins", embed=False):
     def inner(x):
 
         if isinstance(x, Iterable) and not isinstance(x, dict):
@@ -54,6 +62,11 @@ def inspect(msg="ins"):
 
         print(msg, x)
 
+        if embed:
+            import IPython
+
+            IPython.embed()
+
         return x
 
     return inner
diff --git a/test/exploration_tests/partial_response_test.py b/test/exploration_tests/partial_response_test.py
index e0071f2..a6b0adb 100644
--- a/test/exploration_tests/partial_response_test.py
+++ b/test/exploration_tests/partial_response_test.py
@@ -1,15 +1,14 @@
 import logging
-from functools import partial
-from itertools import chain, repeat
+from itertools import repeat
 from operator import methodcaller
 from typing import Iterable
 
 import pytest
 import requests
-from funcy import rcompose, compose, rpartial, identity, lmap, ilen, first, flatten
+from funcy import rcompose, compose, rpartial, identity, ilen, flatten
 
-from pyinfra.rest import pack, unpack, bundle, inspect
-from pyinfra.utils.func import lift, starlift, parallel_map, star, lstarlift
+from pyinfra.rest import pack, inspect
+from pyinfra.utils.func import lift, starlift, parallel_map
 
 logger = logging.getLogger("PIL.PngImagePlugin")
 logger.setLevel(logging.WARNING)
@@ -42,23 +41,8 @@ def post_partial(url, input_data: Iterable[bytes], metadata):
     return input_data_to_result_data(input_data)
 
 
-@pytest.mark.parametrize(
-    "item_type",
-    [
-        # "string",
-        # "image",
-        "pdf"
-    ],
-)
+@pytest.mark.parametrize("batched", [True, False])
+@pytest.mark.parametrize("item_type", ["pdf", "string", "image"])
 def test_sending_partial_request(url, input_data_items, metadata, operation, target_data_items, server_process):
-
-    print()
-    print("exp")
-    print(target_data_items)
-
     output = list(post_partial(f"{url}/process", input_data_items, metadata))
-
-    print()
-    print("out")
-    print(output)
     assert output == target_data_items
diff --git a/test/exploration_tests/pickup_endpoint_test.py b/test/exploration_tests/pickup_endpoint_test.py
index 74149f7..2631ad5 100644
--- a/test/exploration_tests/pickup_endpoint_test.py
+++ b/test/exploration_tests/pickup_endpoint_test.py
@@ -19,8 +19,6 @@ def test_pickup_endpoint(url, server_process, pdf, metadata, operation):
         post,
         rpartial(pack, metadata),
     )(pdf)
-    print(pickup)
 
     # while True:
     #     response = requests.get(f"{url}/{pickup}")
-    #     print(response)
diff --git a/test/fixtures/server.py b/test/fixtures/server.py
index e0a94a7..64f4c2b 100644
--- a/test/fixtures/server.py
+++ b/test/fixtures/server.py
@@ -7,10 +7,10 @@ import fitz
 import pytest
 import requests
 from PIL import Image
-from funcy import retry, compose
+from funcy import retry, compose, flatten
 from waitress import serve
 
-from pyinfra.rest import unpack_op_pack, unpack_batchop_pack
+from pyinfra.rest import unpack_op_pack, unpack_batchop_pack, inspect
 from pyinfra.utils.buffer import bufferize
 from pyinfra.utils.func import llift, starlift
 from test.server import set_up_processing_server