From b2f1df85a36abb7bf634b469c01cf28919026b50 Mon Sep 17 00:00:00 2001
From: "Arif, Maisam" <Maisam.Arif@amd.com>
Date: Tue, 22 Apr 2025 23:52:06 -0500
Subject: [PATCH] Improved import structure (#275)

* Improved import structure

Signed-off-by: Maisam Arif <Maisam.Arif@amd.com>
Change-Id: Id265cbb7dba5ba805b7cf7c353af870fef6cbb4a

---------

Signed-off-by: Maisam Arif <Maisam.Arif@amd.com>
---
 amdsmi_cli/amdsmi_cli.py     | 33 ++++++++++++++++-----------------
 amdsmi_cli/amdsmi_helpers.py |  1 +
 amdsmi_cli/amdsmi_init.py    | 15 +++++++++++----
 3 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/amdsmi_cli/amdsmi_cli.py b/amdsmi_cli/amdsmi_cli.py
index bb935ea0fb..05483a1a4e 100755
--- a/amdsmi_cli/amdsmi_cli.py
+++ b/amdsmi_cli/amdsmi_cli.py
@@ -30,37 +30,36 @@ except ImportError as e:
     logging.debug(f"Unhandled import error: {e}")
     logging.debug("argcomplete module not found. Autocomplete will not work.")
 
-from typing import TYPE_CHECKING
-# only used for type checking
-# pyright trips up and cannot find amdsmi scripts without it
-if TYPE_CHECKING:
-    from amdsmi_commands import AMDSMICommands
-    from amdsmi_parser import AMDSMIParser
-    from amdsmi_logger import AMDSMILogger
-    import amdsmi_cli_exceptions
-    from amdsmi import amdsmi_interface
-    from amdsmi import amdsmi_exception
+# from typing import TYPE_CHECKING
+# # only used for type checking
+# # pyright trips up and cannot find amdsmi scripts without it
+# if TYPE_CHECKING:
+#     from amdsmi_commands import AMDSMICommands
+#     from amdsmi_parser import AMDSMIParser
+#     from amdsmi_logger import AMDSMILogger
+#     import amdsmi_cli_exceptions
+#     from amdsmi import amdsmi_interface
+#     from amdsmi import amdsmi_exception
+
 try:
+    from amdsmi_init import *
     from amdsmi_commands import AMDSMICommands
     from amdsmi_parser import AMDSMIParser
     from amdsmi_logger import AMDSMILogger
     import amdsmi_cli_exceptions
-    from amdsmi import amdsmi_interface
-    from amdsmi import amdsmi_exception
 except ImportError:
     current_path = os.path.dirname(os.path.abspath(__file__))
-    additional_path = f"{current_path}/../libexec/amdsmi_cli"
-    sys.path.append(additional_path)
+    cli_files_path = f"{current_path}/../libexec/amdsmi_cli"
+    sys.path.append(cli_files_path)
     try:
+        from amdsmi_init import *
         from amdsmi_commands import AMDSMICommands
         from amdsmi_parser import AMDSMIParser
         from amdsmi_logger import AMDSMILogger
         import amdsmi_cli_exceptions
-        from amdsmi import amdsmi_interface
-        from amdsmi import amdsmi_exception
     except ImportError as e:
         print(f"Unhandled import error: {e}")
-        print(f"Still couldn't import 'amdsmi related scripts'. Make sure it's installed in {additional_path}")
+        print(f"Unable to import amdsmi_cli files. Check {cli_files_path} if they are present.")
         sys.exit(1)
 
 def _print_error(e, destination):
diff --git a/amdsmi_cli/amdsmi_helpers.py b/amdsmi_cli/amdsmi_helpers.py
index bc2d909025..b7c38fa345 100644
--- a/amdsmi_cli/amdsmi_helpers.py
+++ b/amdsmi_cli/amdsmi_helpers.py
@@ -34,6 +34,7 @@ from enum import Enum
 from pathlib import Path
 from typing import List, Set, Union
 
+# Import amdsmi library
 from amdsmi_init import *
 from BDF import BDF
 
diff --git a/amdsmi_cli/amdsmi_init.py b/amdsmi_cli/amdsmi_init.py
index 6b2e8ced4a..c8c84c12f1 100644
--- a/amdsmi_cli/amdsmi_init.py
+++ b/amdsmi_cli/amdsmi_init.py
@@ -29,11 +29,18 @@ import os
 
 from pathlib import Path
 
-sys.path.append(f"{Path(__file__).resolve().parent}/../../share/amd_smi")
-
+current_path = os.path.dirname(os.path.abspath(__file__))
+python_lib_path = f"{current_path}/../../share/amd_smi"
+sys.path.append(python_lib_path)
 # If the python library is installed, it will overwrite the path above
-from amdsmi import amdsmi_interface
-from amdsmi import amdsmi_exception
+
+try:
+    from amdsmi import amdsmi_interface, amdsmi_exception
+except ImportError as e:
+    print(f"Unhandled import error: {e}")
+    print("Failed to import the amdsmi Python library. Ensure it is installed in Python.")
+    print(f"Alternatively, verify that the library is in the path:\n{python_lib_path}")
+    sys.exit(1)
 
 # Using basic python logging for user errors and development
 logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.ERROR) # User level logging