use FlowerCallable instead of App

src/py/flwr/client/flower.py

@@ -19,9 +19,8 @@
 from typing import List, Optional, cast
 
 from flwr.client.message_handler.message_handler import handle
-from flwr.client.middleware.typing import Layer
 from flwr.client.middleware.utils import make_app
-from flwr.client.typing import Bwd, ClientFn, Fwd
+from flwr.client.typing import Bwd, ClientFn, Fwd, Layer
 
 
 class Flower:

src/py/flwr/client/middleware/__init__.py

@@ -15,11 +15,8 @@
 """Middleware layers."""
 
 
-from .typing import App, Layer
 from .utils import make_app
 
 __all__ = [
-    "App",
-    "Layer",
     "make_app",
 ]

src/py/flwr/client/middleware/utils.py

@@ -16,15 +16,13 @@
 
 from typing import List
 
-from flwr.client.typing import Bwd, Fwd
+from flwr.client.typing import Bwd, FlowerCallable, Fwd, Layer
 
-from .typing import App, Layer
 
-
-def make_app(app: App, middleware_layers: List[Layer]) -> App:
+def make_app(app: FlowerCallable, middleware_layers: List[Layer]) -> FlowerCallable:
     """."""
 
-    def wrap_app(_app: App, _layer: Layer) -> App:
+    def wrap_app(_app: FlowerCallable, _layer: Layer) -> FlowerCallable:
         def new_app(fwd: Fwd) -> Bwd:
             return _layer(fwd, _app)
 

src/py/flwr/client/middleware/utils_test.py

@@ -18,18 +18,17 @@
 import unittest
 from typing import List
 
-from flwr.client.typing import Bwd, Fwd
+from flwr.client.typing import Bwd, FlowerCallable, Fwd, Layer
 from flwr.client.workload_state import WorkloadState
 from flwr.proto.task_pb2 import TaskIns, TaskRes
 
-from .typing import App, Layer
 from .utils import make_app
 
 
 def make_mock_middleware(name: str, footprint: List[str]) -> Layer:
     """Make a mock middleware layer."""
 
-    def middleware(fwd: Fwd, app: App) -> Bwd:
+    def middleware(fwd: Fwd, app: FlowerCallable) -> Bwd:
         footprint.append(name)
         fwd.task_ins.task_id += f"{name}"
         bwd = app(fwd)
@@ -40,7 +39,7 @@ def middleware(fwd: Fwd, app: App) -> Bwd:
     return middleware
 
 
-def make_mock_app(name: str, footprint: List[str]) -> App:
+def make_mock_app(name: str, footprint: List[str]) -> FlowerCallable:
     """Make a mock app."""
 
     def app(fwd: Fwd) -> Bwd:
@@ -83,7 +82,7 @@ def test_filter(self) -> None:
         mock_app = make_mock_app("app", footprint)
         task_ins = TaskIns()
 
-        def filter_layer(fwd: Fwd, _: App) -> Bwd:
+        def filter_layer(fwd: Fwd, _: FlowerCallable) -> Bwd:
             footprint.append("filter")
             fwd.task_ins.task_id += "filter"
             # Skip calling app

src/py/flwr/client/typing.py

@@ -41,3 +41,4 @@ class Bwd:
 
 FlowerCallable = Callable[[Fwd], Bwd]
 ClientFn = Callable[[str], Client]
+Layer = Callable[[Fwd, FlowerCallable], Bwd]