From 094e92519f1025a8bd1bdc5885a06dce83616349 Mon Sep 17 00:00:00 2001
From: Heng Pan <134433891+panh99@users.noreply.github.com>
Date: Thu, 30 Nov 2023 14:52:37 +0000
Subject: [PATCH] Move definitions of `Fwd` and `Bwd` to `client/typing.py`
 (#2669)

---
 examples/vertical-fl/simulation.py |  2 +-
 src/py/flwr/client/app.py          |  4 ++--
 src/py/flwr/client/flower.py       | 26 ++------------------------
 src/py/flwr/client/typing.py       | 22 ++++++++++++++++++++++
 src/py/flwr/flower/__init__.py     |  4 ++--
 5 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/examples/vertical-fl/simulation.py b/examples/vertical-fl/simulation.py
index 095ec6e0c7ea..f4befc3a073e 100644
--- a/examples/vertical-fl/simulation.py
+++ b/examples/vertical-fl/simulation.py
@@ -22,4 +22,4 @@ def client_fn(cid):
 
 results_dir = Path("_static/results")
 results_dir.mkdir(exist_ok=True)
-np.save(str(results_dir/"hist.npy"), hist)
+np.save(str(results_dir / "hist.npy"), hist)
diff --git a/src/py/flwr/client/app.py b/src/py/flwr/client/app.py
index 81bbee148c95..7ce7d51d3d4b 100644
--- a/src/py/flwr/client/app.py
+++ b/src/py/flwr/client/app.py
@@ -23,8 +23,8 @@
 from typing import Callable, ContextManager, Optional, Tuple, Union
 
 from flwr.client.client import Client
-from flwr.client.flower import Bwd, Flower, Fwd
-from flwr.client.typing import ClientFn
+from flwr.client.flower import Flower
+from flwr.client.typing import Bwd, ClientFn, Fwd
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH, EventType, event
 from flwr.common.address import parse_address
 from flwr.common.constant import (
diff --git a/src/py/flwr/client/flower.py b/src/py/flwr/client/flower.py
index f0d4ce122524..5b083ee11b9f 100644
--- a/src/py/flwr/client/flower.py
+++ b/src/py/flwr/client/flower.py
@@ -16,32 +16,10 @@
 
 
 import importlib
-from dataclasses import dataclass
-from typing import Callable, cast
+from typing import cast
 
 from flwr.client.message_handler.message_handler import handle
-from flwr.client.typing import ClientFn
-from flwr.client.workload_state import WorkloadState
-from flwr.proto.task_pb2 import TaskIns, TaskRes
-
-
-@dataclass
-class Fwd:
-    """."""
-
-    task_ins: TaskIns
-    state: WorkloadState
-
-
-@dataclass
-class Bwd:
-    """."""
-
-    task_res: TaskRes
-    state: WorkloadState
-
-
-FlowerCallable = Callable[[Fwd], Bwd]
+from flwr.client.typing import Bwd, ClientFn, Fwd
 
 
 class Flower:
diff --git a/src/py/flwr/client/typing.py b/src/py/flwr/client/typing.py
index 7ee6f069768c..2c1f7506592c 100644
--- a/src/py/flwr/client/typing.py
+++ b/src/py/flwr/client/typing.py
@@ -14,8 +14,30 @@
 # ==============================================================================
 """Custom types for Flower clients."""
 
+from dataclasses import dataclass
 from typing import Callable
 
+from flwr.client.workload_state import WorkloadState
+from flwr.proto.task_pb2 import TaskIns, TaskRes
+
 from .client import Client as Client
 
+
+@dataclass
+class Fwd:
+    """."""
+
+    task_ins: TaskIns
+    state: WorkloadState
+
+
+@dataclass
+class Bwd:
+    """."""
+
+    task_res: TaskRes
+    state: WorkloadState
+
+
+FlowerCallable = Callable[[Fwd], Bwd]
 ClientFn = Callable[[str], Client]
diff --git a/src/py/flwr/flower/__init__.py b/src/py/flwr/flower/__init__.py
index 090c78062d02..892a7ce5afdc 100644
--- a/src/py/flwr/flower/__init__.py
+++ b/src/py/flwr/flower/__init__.py
@@ -15,9 +15,9 @@
 """Flower callable package."""
 
 
-from flwr.client.flower import Bwd as Bwd
 from flwr.client.flower import Flower as Flower
-from flwr.client.flower import Fwd as Fwd
+from flwr.client.typing import Bwd as Bwd
+from flwr.client.typing import Fwd as Fwd
 
 __all__ = [
     "Flower",