From 23e3ed011d240e5e8e996d3c321a102e976841db Mon Sep 17 00:00:00 2001
From: Pierre Fayolle <pierre.fayolle@vossloh.com>
Date: Tue, 16 Jan 2024 10:51:37 +0100
Subject: [PATCH] Try relative imports of tahu files if global import fails

Removes the need to add tahu to the PYTHONPATH
---
 enki.py       |  5 ++++-
 shell.py      | 12 +++++++++---
 sp_helpers.py |  5 ++++-
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/enki.py b/enki.py
index f6cd36c..55724d7 100755
--- a/enki.py
+++ b/enki.py
@@ -5,7 +5,10 @@
 import argparse
 import signal
 
-from sparkplug_b_pb2 import Payload
+try:
+    from sparkplug_b_pb2 import Payload
+except ImportError:
+    from tahu.python.core.sparkplug_b_pb2 import Payload
 
 from mqtt_if import MQTTInterface
 from sp_topic import SPTopic
diff --git a/shell.py b/shell.py
index 0111104..d7ec881 100644
--- a/shell.py
+++ b/shell.py
@@ -5,8 +5,14 @@
 
 import cmd2
 
-import sparkplug_b_pb2
-from sparkplug_b import MetricDataType, addMetric, addNullMetric, initDatasetMetric
+try:
+    from sparkplug_b_pb2 import Payload
+    from sparkplug_b import MetricDataType, addMetric
+    from sparkplug_b import addNullMetric, initDatasetMetric
+except ImportError:
+    from tahu.python.core.sparkplug_b_pb2 import Payload
+    from tahu.python.core.sparkplug_b import MetricDataType, addMetric
+    from tahu.python.core.sparkplug_b import addNullMetric, initDatasetMetric
 
 import sp_helpers
 from sp_helpers import MsgType
@@ -433,7 +439,7 @@ def do_send(self, args):
 
         msg_type = MsgType.CMD if args.msg_type == "CMD" else MsgType.DATA
         topic = sp_dev.get_msg_topic(msg_type)
-        payload = sparkplug_b_pb2.Payload()
+        payload = Payload()
         new_metric = True
         while new_metric:
             metric = self.choose_metric(sp_dev)
diff --git a/sp_helpers.py b/sp_helpers.py
index 9e588bb..3eb15b3 100644
--- a/sp_helpers.py
+++ b/sp_helpers.py
@@ -4,7 +4,10 @@
 from typing import Optional
 from enum import Enum
 
-from sparkplug_b import MetricDataType
+try:
+    from sparkplug_b import MetricDataType
+except ImportError:
+    from tahu.python.core.sparkplug_b import MetricDataType
 
 
 # Sparkplug namespace used