From 9a292d61cd4555733c3ed6e3a8ca016b17e062c3 Mon Sep 17 00:00:00 2001
From: John Chilton <jmchilton@gmail.com>
Date: Thu, 31 Aug 2023 10:52:01 -0400
Subject: [PATCH] More typing?

---
 lib/galaxy/tool_shed/galaxy_install/client.py   | 8 ++++++--
 lib/galaxy/tool_shed/unittest_utils/__init__.py | 9 ++++++---
 lib/galaxy/tools/__init__.py                    | 8 ++++----
 lib/galaxy/tools/data_manager/manager.py        | 2 +-
 4 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/lib/galaxy/tool_shed/galaxy_install/client.py b/lib/galaxy/tool_shed/galaxy_install/client.py
index c537b9d1def0..92125ba6d78d 100644
--- a/lib/galaxy/tool_shed/galaxy_install/client.py
+++ b/lib/galaxy/tool_shed/galaxy_install/client.py
@@ -1,6 +1,7 @@
 import threading
 from typing import (
     Any,
+    Dict,
     Generic,
     List,
     Optional,
@@ -15,7 +16,10 @@
 from galaxy.model.tool_shed_install import HasToolBox
 from galaxy.security.idencoding import IdEncodingHelper
 from galaxy.tool_shed.cache import ToolShedRepositoryCache
-from galaxy.tool_util.data import ToolDataTableManager
+from galaxy.tool_util.data import (
+    OutputDataset,
+    ToolDataTableManager,
+)
 from galaxy.tool_util.toolbox.base import AbstractToolBox
 
 if TYPE_CHECKING:
@@ -29,7 +33,7 @@ class DataManagerInterface(Protocol):
     def process_result(self, out_data):
         ...
 
-    def write_bundle(self, out) -> None:
+    def write_bundle(self, out: Dict[str, OutputDataset]) -> Dict[str, OutputDataset]:
         ...
 
 
diff --git a/lib/galaxy/tool_shed/unittest_utils/__init__.py b/lib/galaxy/tool_shed/unittest_utils/__init__.py
index 61062d76164c..2785fbca8857 100644
--- a/lib/galaxy/tool_shed/unittest_utils/__init__.py
+++ b/lib/galaxy/tool_shed/unittest_utils/__init__.py
@@ -23,7 +23,10 @@
     InstallationTarget,
 )
 from galaxy.tool_shed.util.repository_util import get_installed_repository
-from galaxy.tool_util.data import ToolDataTableManager
+from galaxy.tool_util.data import (
+    OutputDataset,
+    ToolDataTableManager,
+)
 from galaxy.tool_util.loader_directory import looks_like_a_tool
 from galaxy.tool_util.toolbox.base import (
     AbstractToolBox,
@@ -134,8 +137,8 @@ class DummyDataManager(DataManagerInterface):
     def process_result(self, out_data):
         return None
 
-    def write_bundle(self, out) -> None:
-        return None
+    def write_bundle(self, out) -> Dict[str, OutputDataset]:
+        return {}
 
 
 class StandaloneDataManagers(DataManagersInterface):
diff --git a/lib/galaxy/tools/__init__.py b/lib/galaxy/tools/__init__.py
index 405e80035768..3d03e0348d51 100644
--- a/lib/galaxy/tools/__init__.py
+++ b/lib/galaxy/tools/__init__.py
@@ -3108,10 +3108,10 @@ def exec_after_process(self, app, inp_data, out_data, param_dict, job=None, fina
             pass
         elif data_manager_mode == "bundle":
             for bundle_path, dataset in data_manager.write_bundle(out_data).items():
-                dataset = cast(model.HistoryDatasetAssociation, dataset)
-                dataset.dataset.object_store.update_from_file(
-                    dataset.dataset,
-                    extra_dir=dataset.dataset.extra_files_path_name,
+                hda = cast(model.HistoryDatasetAssociation, dataset)
+                hda.dataset.object_store.update_from_file(
+                    hda.dataset,
+                    extra_dir=hda.dataset.extra_files_path_name,
                     file_name=bundle_path,
                     alt_name=os.path.basename(bundle_path),
                     create=True,
diff --git a/lib/galaxy/tools/data_manager/manager.py b/lib/galaxy/tools/data_manager/manager.py
index 622e7189d20f..52f79ca54250 100644
--- a/lib/galaxy/tools/data_manager/manager.py
+++ b/lib/galaxy/tools/data_manager/manager.py
@@ -244,7 +244,7 @@ def process_result(self, out_data: Dict[str, OutputDataset]) -> None:
     def write_bundle(
         self,
         out_data: Dict[str, OutputDataset],
-    ):
+    ) -> Dict[str, OutputDataset]:
         tool_data_tables = self.data_managers.app.tool_data_tables
         return tool_data_tables.write_bundle(
             out_data,