From 087b5af929624d9bd1d5b0320bb32bac74a70f66 Mon Sep 17 00:00:00 2001
From: Matthias Bisping <matthias.bisping@iqser.com>
Date: Thu, 28 Apr 2022 18:28:35 +0200
Subject: [PATCH] refactoring

---
 pyinfra/rest.py                               |  5 ++-
 pyinfra/utils/func.py                         | 30 +++++++++++++++++
 .../partial_response_test.py                  | 32 ++-----------------
 3 files changed, 35 insertions(+), 32 deletions(-)
 create mode 100644 pyinfra/utils/func.py

diff --git a/pyinfra/rest.py b/pyinfra/rest.py
index 24bff6e..e83d087 100644
--- a/pyinfra/rest.py
+++ b/pyinfra/rest.py
@@ -1,9 +1,8 @@
-import json
-from _operator import itemgetter
 from collections import deque
+from operator import itemgetter
 
 import flask
-from funcy import repeatedly, lmap, compose
+from funcy import repeatedly, lmap, compose, zipdict, rcompose, identity
 
 from test.utils.server import bytes_to_string, string_to_bytes
 
diff --git a/pyinfra/utils/func.py b/pyinfra/utils/func.py
new file mode 100644
index 0000000..059bfb3
--- /dev/null
+++ b/pyinfra/utils/func.py
@@ -0,0 +1,30 @@
+from itertools import starmap, tee
+
+from funcy import curry, compose
+
+
+def lift(fn):
+    return curry(map)(fn)
+
+
+def starlift(fn):
+    return curry(starmap)(fn)
+
+
+def parallel(*fs):
+    return lambda *args: (f(a) for f, a in zip(fs, args))
+
+
+def star(f):
+    return lambda x: f(*x)
+
+
+def duplicate_stream_and_apply(f1, f2):
+    return compose(star(parallel(f1, f2)), tee)
+
+
+def parallel_map(f1, f2):
+    """Applies functions to a stream in parallel and yields a stream of tuples:
+    parallel_map :: a -> b, a -> c -> [a] -> [(b, c)]
+    """
+    return compose(star(zip), duplicate_stream_and_apply(f1, f2))
diff --git a/test/exploration_tests/partial_response_test.py b/test/exploration_tests/partial_response_test.py
index 3c881c1..e320532 100644
--- a/test/exploration_tests/partial_response_test.py
+++ b/test/exploration_tests/partial_response_test.py
@@ -1,47 +1,21 @@
 import logging
-from itertools import chain, starmap, tee
+from itertools import chain
 from operator import methodcaller, itemgetter
 from typing import Iterable
 
 import pytest
 import requests
-from funcy import curry, rcompose, compose, lmap, rpartial, identity
+from funcy import rcompose, compose, lmap, rpartial, identity
 from more_itertools import peekable
 
 from pyinfra.rest import pack
+from pyinfra.utils.func import lift, starlift, parallel_map
 from test.utils.server import string_to_bytes
 
 logger = logging.getLogger("PIL.PngImagePlugin")
 logger.setLevel(logging.INFO)
 
 
-def lift(fn):
-    return curry(map)(fn)
-
-
-def starlift(fn):
-    return curry(starmap)(fn)
-
-
-def parallel(*fs):
-    return lambda *args: (f(a) for f, a in zip(fs, args))
-
-
-def star(f):
-    return lambda x: f(*x)
-
-
-def duplicate_stream_and_apply(f1, f2):
-    return compose(star(parallel(f1, f2)), tee)
-
-
-def parallel_map(f1, f2):
-    """Applies functions to a stream in parallel and yields a stream of tuples:
-    parallel_map :: a -> b, a -> c -> [a] -> [(b, c)]
-    """
-    return compose(star(zip), duplicate_stream_and_apply(f1, f2))
-
-
 def post_partial(url, input_data: Iterable[bytes], metadata):
     def send(method, data):
         return method(url, json=data)