Add support for user-defined macros

Signed-off-by: John Craig <[email protected]>

Author: John Craig

examples/macros/logon.py

@@ -1,8 +1,7 @@
 from tnz.ati import ati
 
-def logon(zti, arg):
+def do_logon(zti, arg):
     logon_setup()
-
     ati.wait(lambda: ati.scrhas("Enter:  "))
     ati.send("app1 userid[enter]")
     ati.wait(lambda: ati.scrhas("Password  ===> "))
@@ -21,11 +20,11 @@ def logon(zti, arg):
     # disconnect from host
     ati.drop("SESSION")
 
-def logon_helper():
+def logon_setup():
     ati.set("TRACE", "ALL")
     ati.set("LOGDEST", "example.log")
 
     ati.set("ONERROR", "1")
     ati.set("DISPLAY", "HOST")
     ati.set("SESSION_HOST", "mvs1")
-    ati.set("SESSION", "A")
+    ati.set("SESSION", "A")

tnz/zti.py

@@ -120,11 +120,6 @@ def __init__(self, stdin=None, stdout=None):
         if os.getenv("ZTI_AUTOSIZE"):
             self.autosize = True
 
-        self.macros_dir = os.getenv("ZTI_MACROS_DIR")
-        if self.macros_dir is None or not os.path.isdir(
-           self.macros_dir):
-            self.macros_dir = os.path.expanduser("~/.zti-macros.d")
-
         self.pend_intro = None
         self.__dirty_ranges = []
         self.single_session = False
@@ -1861,45 +1856,6 @@ def help_plugin(entry=entry):
 
         self.plugins = " ".join(plugins)
 
-    def __register_macros(self):
-        import importlib.util
-        import sys
-        import types
-
-        if not os.path.isdir(self.macros_dir):
-            return
-
-        for macro_file in os.listdir(self.macros_dir):
-            if not macro_file.endswith(".py"):
-                continue
-
-            if len(macro_file.split('.')) > 1:
-                continue
-
-            macro_name = macro_file.split('.')[0]
-
-            # Ignore macros which already exist
-            if f"do_{macro_name}" in self.get_names():
-                continue
-
-            macro_path = os.path.join(self.macros_dir, macro_file)
-
-            # Import the user macro as a module
-            macro_spec = importlib.util.spec_from_file_location(
-                f"module.{macro_name}", macro_path)
-            macro_module = importlib.util.module_from_spec(macro_spec)
-            sys.modules[f"module.{macro_name}"] = macro_module
-            macro_spec.loader.exec_module(macro_module)
-
-            # Find the function
-            if hasattr(macro_module, macro_name):
-                # Create a new bound method for the `Zti` class for this
-                # function
-                setattr(Zti, f"do_{macro_name}",
-                        types.MethodType(
-                            getattr(macro_module, macro_name),
-                            self))
-
     def __key_data(self, tns, data):
         try:
             tns.key_data(data, zti=self)
@@ -2751,6 +2707,71 @@ def __refresh(self):
 
         self.stdscr.refresh(_win_callback=self.__set_event_fn())
 
+    def __register_macros(self):
+        import importlib.util
+        import sys
+        import types
+
+        macros_dir = os.getenv("ZTI_MACROS_DIR")
+        if macros_dir is None:
+            macros_dir = os.path.expanduser("~/.zti-mac")
+
+        if not os.path.isdir(macros_dir):
+            _logger.exception(f"{macros_dir} is not a directory")
+            return
+
+        for macro_file in os.listdir(macros_dir):
+            if not macro_file.endswith(".py"):
+                continue
+
+            if len(macro_file.split('.')) > 1:
+                continue
+
+            macro_name = macro_file.split('.')[0]
+
+            # Ignore macros with uppercase letters or
+            # spaces
+            uppercase = [l for l in macro_name if l.isupper()]
+            if ' ' in macro_name or len(uppercase) != 0:
+                continue
+
+            # Ignore macros which already exist
+            if f"do_{macro_name}" in self.get_names():
+                _logger.warning(f"Function with name do_{macro_name} already exists, macro registration failed")
+                continue
+
+            macro_path = os.path.join(macros_dir, macro_file)
+
+            # Import the user macro as a module
+            macro_spec = importlib.util.spec_from_file_location(
+                f"module.{macro_name}", macro_path)
+            macro_module = importlib.util.module_from_spec(macro_spec)
+            sys.modules[f"module.{macro_name}"] = macro_module
+            macro_spec.loader.exec_module(macro_module)
+
+            # Find the function
+            if hasattr(macro_module, f"do_{macro_name}"):
+                def do_macro(zti, arg):
+                    self.__bg_wait_end()
+                    macro_func = getattr(macro_module, f"do_{macro_name}")
+                    macro_func(zti, arg)
+
+                # Create a new bound method for the `Zti` class for this
+                # function
+                setattr(Zti, f"do_{macro_name}",
+                        types.MethodType(
+                            do_macro,
+                            self))
+
+                # Check if a corresponding help function exists
+                if hasattr(macro_module, f"help_{macro_name}"):
+                    # Create a new bound method for the `Zti` class for this
+                    # function
+                    setattr(Zti, f"help_{macro_name}",
+                            types.MethodType(
+                                getattr(macro_module, f"help_{macro_name}"),
+                                self))
+
     def __scale_size(self, maxrow, maxcol):
         arows, acols = self.autosize
         aspect1 = arows / acols