From 3db077d15019486131ad3d949b6b53131e82d1cf Mon Sep 17 00:00:00 2001
From: Jonathon Broughton <jonathon@stardotbmp.com>
Date: Mon, 24 Mar 2025 13:38:56 +0000
Subject: [PATCH] Refactor imports for clarity and organisation

- Removed unnecessary imports from the main module.
- Consolidated import statements for `FunctionInputs` and `automate_function`.
- Updated test file to reflect new import paths.
---
 main.py                | 11 +++--------
 tests/test_function.py |  3 +--
 2 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/main.py b/main.py
index ff6c9c8..f2471e9 100644
--- a/main.py
+++ b/main.py
@@ -1,15 +1,10 @@
-from speckle_automate import (
-    execute_automate_function,
-)
-
-
-from src.function import automate_function
-from src.inputs import FunctionInputs
+from speckle_automate import execute_automate_function
 
+from data_shield import FunctionInputs
+from data_shield import automate_function
 
 # make sure to call the function with the executor
 if __name__ == "__main__":
     # NOTE: always pass in the automate function by its reference; do not invoke it!
 
-
     execute_automate_function(automate_function, FunctionInputs)
diff --git a/tests/test_function.py b/tests/test_function.py
index e577732..dc7fb0d 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -7,8 +7,7 @@ from speckle_automate import (
 )
 from speckle_automate.fixtures import *  # noqa: F401, F403
 
-from src.function import automate_function
-from src.inputs import FunctionInputs, SanitizationMode
+from data_shield import automate_function, FunctionInputs, SanitizationMode
 
 
 class TestFunction: