Add to_play query parameter

codeguru42 · codeguru42 · commit 9d2e57b63802 · 2023-07-15T15:08:22.000-06:00
diff --git a/main.py b/main.py
@@ -1,4 +1,5 @@
 import io
+from enum import Enum
 from pathlib import Path
 from typing import Annotated
 
@@ -11,15 +12,23 @@
 app = FastAPI()
 
 
+class Color(str, Enum):
+    BLACK = "B"
+    WHITE = "W"
+
+
 @app.get("/health_check/")
 def health_check():
     return {"message": "Healthy!"}
 
 
 @app.post("/capture/")
-def capture(image: Annotated[UploadFile, File()]):
+def capture(
+    image: Annotated[UploadFile, File()],
+    to_play: Annotated[Color, Form()] = Color.BLACK,
+):
     output_file = io.StringIO()
-    process_image(image.file, output_file)
+    process_image(image.file, output_file, to_play)
     output_file.seek(0)
     return {
         "sgf": output_file.read(),
@@ -29,10 +38,14 @@ def capture(image: Annotated[UploadFile, File()]):
 
 @app.post("/capture_async/")
 def capture_async(
-    image: Annotated[UploadFile, File()], fcm_registration_token: Annotated[str, Form()]
+    image: Annotated[UploadFile, File()],
+    fcm_registration_token: Annotated[str, Form()],
+    to_play: Annotated[Color, Form()] = Color.BLACK,
 ):
     filename = Path(image.filename)
     output_path = settings.IMAGES_DIR / filename
     with output_path.open("wb") as output_file:
         output_file.write(image.file.read())
-    process_image_task.delay(str(output_path.absolute()), fcm_registration_token)
+    process_image_task.delay(
+        str(output_path.absolute()), fcm_registration_token, to_play
+    )
