sync

pyinfra/rest.py

@@ -0,0 +1,9 @@
+import json
+
+from test.utils.server import bytes_to_string
+
+
+def pack(data: bytes, metadata: dict):
+    package = {"data": bytes_to_string(data), "metadata": metadata}
+    package = json.dumps(package).encode()
+    return package

requirements.txt

@@ -13,3 +13,4 @@ tqdm==4.62.3
 pytest~=7.0.1
 funcy==1.17
 fpdf==1.7.2
+PyMuPDF==1.19.6

test/exploration_tests/partial_response_test.py

@@ -1,4 +1,3 @@
-import json
 import logging
 from itertools import chain
 from operator import methodcaller, itemgetter
@@ -6,9 +5,10 @@ from typing import Iterable
 
 import pytest
 import requests
-from funcy import curry, rcompose, compose, lmap
+from funcy import curry, rcompose, compose, lmap, rpartial
 
-from test.utils.server import bytes_to_string, string_to_bytes
+from pyinfra.rest import pack
+from test.utils.server import string_to_bytes
 
 logger = logging.getLogger("PIL.PngImagePlugin")
 logger.setLevel(logging.INFO)
@@ -19,15 +19,10 @@ def lift(fn):
 
 
 def post_partial(url, input_data: Iterable[bytes], metadata):
-    def pack(data: bytes):
-        package = {"data": bytes_to_string(data), "metadata": metadata}
-        package = json.dumps(package).encode()
-        return package
-
     def post(data):
         return requests.post(url, data=data)
 
-    pack_data_and_metadata_for_rest_transfer = lift(pack)
+    pack_data_and_metadata_for_rest_transfer = lift(rpartial(pack, metadata))
     send_packages_to_analyzer_and_receive_responses = lift(post)
     extract_data_from_responses = lift(compose(itemgetter("data"), methodcaller("json")))
     flatten_buffered_payloads = chain.from_iterable
@@ -50,8 +45,3 @@ def test_sending_partial_request(url, data_items, metadata, operation, server_pr
     output = list(post_partial(f"{url}/process", data_items, metadata))
 
     assert output == expected
-
-
-@pytest.fixture
-def metadata():
-    return {"idx": [1, 100, 101], "path": "asd/asd"}

test/exploration_tests/pickup_endpoint_test.py

@@ -0,0 +1,25 @@
+import json
+from operator import itemgetter, methodcaller
+
+import pytest
+import requests
+from funcy import compose, rpartial
+
+from pyinfra.rest import pack
+
+
+@pytest.mark.parametrize("item_type", ["pdf"])
+def test_sending_partial_request(url, server_process, pdf, metadata, operation):
+
+    def post(data):
+        return requests.post(f"{url}/submit", data=data)
+
+    pickup = compose(itemgetter("pickup_endpoint"), methodcaller("json"), post, rpartial(pack, metadata))(pdf)
+    print(pickup)
+
+    while True:
+        response = requests.get(f"{url}/{pickup}")
+        print(response)
+
+
+

test/fixtures/input.py

@@ -1,6 +1,3 @@
-from itertools import repeat
-
-
 import numpy as np
 import pytest
 from PIL import Image
@@ -37,3 +34,8 @@ def input_batch():
 
 def images():
     return lmap(compose(image_to_bytes, array_to_image), input_batch())
+
+
+@pytest.fixture
+def metadata():
+    return {"idx": [1, 100, 101], "path": "asd/asd"}

test/fixtures/server.py

@@ -5,6 +5,7 @@ from collections import deque
 from multiprocessing import Process
 from operator import itemgetter
 
+import fitz
 import flask
 import pytest
 import requests
@@ -54,10 +55,16 @@ def operation(item_type):
         im = Image.open(io.BytesIO(im))
         return image_to_bytes(im.rotate(90))
 
+    def stream_pages(pdf: bytes):
+        for page in fitz.open(stream=pdf):
+            yield page.get_pixmap().tobytes("png")
+
     if item_type == "string":
         return lambda s: s.decode().upper().encode()
     elif item_type == "image":
         return rotate
+    elif item_type == "pdf":
+        return stream_pages
     else:
         raise ValueError(f"No operation specified for item type {item_type}")
 

test/server.py

@@ -1,8 +1,10 @@
 from flask import Flask, jsonify, request
+from more_itertools import peekable
 
 
 def set_up_processing_server(process_fn):
     app = Flask(__name__)
+    response_payload_iter = None
 
     @app.route("/ready", methods=["GET"])
     def ready():
@@ -15,4 +17,21 @@ def set_up_processing_server(process_fn):
         response_payload = process_fn(request)
         return jsonify(response_payload)
 
+    @app.route("/submit", methods=["POST"])
+    def submit():
+        nonlocal response_payload_iter
+        response_payload_iter = peekable(iter(process_fn(request)))
+        return jsonify({"pickup_endpoint": "pickup"})
+
+    @app.route("/pickup", methods=["GET"])
+    def pickup():
+
+        response_payload = next(response_payload_iter)
+        print(response_payload)
+
+        resp = jsonify({"a": 1})
+        resp.status_code = 200
+
+        return resp
+
     return app