Refactor functions and add error handling

- run_cmd: remove verbose, add shell arg; call methods directly w/stdin and shell added
- tfa: remove non-existent flag
- get_ip: call subprocess popen and run directly for custom parsing of tf state
- add_ip / rm_ip: inline edit via `re.sub` method
- main: remove annotation, remove prompt in rm_ip() call; added newline re: pep8 before dunder method call

Author: pythoninthegrass

main.py

@@ -6,8 +6,7 @@
 import typer
 from decouple import config
 from pathlib import Path
-from sh import jq
-# from sh import terraform
+from subprocess import Popen, PIPE, STDOUT
 from typing import List, Optional
 from typing_extensions import Annotated
 
@@ -19,7 +18,8 @@
 tenant_id       = config('TENANT_ID')
 
 tf = "terraform"
-tf_dir = str(Path("terraform").resolve())
+tld = "terraform"
+tf_dir = str(Path(f"{tld}/linux").resolve())
 ch_dir = f"-chdir={tf_dir}"
 
 ans_dir = str(Path("ansible").resolve())
@@ -29,34 +29,30 @@
 )
 
 
-def run_cmd(args, verbose=True):
+def run_cmd(args, shell=False):
     """Run command, transfer stdout/stderr back into this process's stdout/stderr"""
     print("Running command: %s" % " ".join(args))
-    if verbose:
-        p = subprocess.Popen(
-            args,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT)
-        out = []
-        with p.stdout:
-            for line in iter(p.stdout.readline, b''):
-                out += [line.decode('utf-8').rstrip()]
-        p.wait()
-        return "\n".join([str(i) for i in out])
-    else:
-        p = subprocess.Popen(
-            args,
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.STDOUT)
-        p.wait()
-        return p.returncode
+    p = Popen(
+        args,
+        shell=shell,
+        stdin=PIPE,
+        stdout=PIPE,
+        stderr=STDOUT)
+    out = []
+    with p.stdout:
+        for line in iter(p.stdout.readline, b''):
+            out += [line.decode('utf-8').rstrip()]
+    p.wait()
+    stdout = "\n".join([str(i) for i in out])
+    print(stdout)
+    return stdout
 
 
 @app.command("hello")
 def hello_world():
     """Debugging command"""
     cmd = ["echo", "hello world"]
-    return print(f"Command exited with code: {run_cmd(cmd, verbose=False)}")
+    return run_cmd(cmd)
 
 
 @app.command()
@@ -79,7 +75,7 @@ def tfa():
     if not Path(f"{tf_dir}/tfplan").exists():
         print("tfplan does not exist, running plan first")
         tfp()
-    cmd = [tf, ch_dir, "apply", "tfplan", "-auto-approve"]
+    cmd = [tf, ch_dir, "apply", "tfplan"]
     return run_cmd(cmd)
 
 
@@ -93,49 +89,67 @@ def tfs():
 @app.command()
 def get_ip():
     """Get Azure instance IP"""
-    # tf_state = terraform.show("-json", _cwd=tf_dir)
-    tf_state = tfs()
-    az_ip = jq(
-        "-r",
-        '.values.root_module.resources[] | select(.address == "azurerm_linux_virtual_machine.my_terraform_vm").values.public_ip_address',
-        _in=tf_state,
+    tf_state = Popen(
+        [tf, ch_dir, "show", "-json"],
+        stdin=PIPE,
+        stdout=PIPE,
     )
+
+    out = []
+    with tf_state.stdout:
+        for line in iter(tf_state.stdout.readline, b''):
+            out += [line.decode('utf-8').rstrip()]
+
+    az_ip = subprocess.run(
+        ["jq", "-r", '.values.root_module.resources[] | select(.address == "azurerm_linux_virtual_machine.my_terraform_vm").values.public_ip_address'],
+        input="\n".join([str(i) for i in out]),
+        capture_output=True,
+        text=True,
+    ).stdout.strip()
+
     print(az_ip)
-    return az_ip.strip()
+    return az_ip
 
 
-# TODO: debug args + match statement in main
 @app.command()
-def add_ip(*args: List[str]):
+def add_ip(*args):
     """Add Azure instance IP to Ansible inventory"""
     if args:
         az_ip = args[0]
     else:
         az_ip = get_ip()
     hosts = Path(f"{ans_dir}/hosts").read_text()
     if az_ip not in hosts:
-        hosts = hosts.replace("[azure]", f"[azure]\n{az_ip}")
+        hosts = re.sub(r"\[azure\]\n", f"[azure]\n{az_ip}\n", hosts)
         Path(f"{ans_dir}/hosts").write_text(hosts)
         return print(Path(f"{ans_dir}/hosts").read_text())
     else:
         print(f"{az_ip} already in hosts file")
 
 
 @app.command()
-def rm_ip(ip_address: str):
+def rm_ip(*args):
     """Remove Azure instance IP from Ansible inventory"""
     hosts = Path(f"{ans_dir}/hosts").read_text()
+    if not args:
+        ip_address = get_ip()
+    elif args:
+        ip_address = args[0]
+    else:
+        print("No IP address specified")
+        raise typer.Abort()
+
     if ip_address in hosts:
-        # delete line with list comprehension
-        hosts = "\n".join([i for i in hosts.split("\n") if i != ip_address])
+        hosts = re.sub(rf"{ip_address}\n", "", hosts)
         Path(f"{ans_dir}/hosts").write_text(hosts)
         return print(Path(f"{ans_dir}/hosts").read_text())
     else:
         print(f"{ip_address} not in hosts file")
 
 
+# TODO: match multiple cases (case guard?)
 @app.callback(invoke_without_command=True)
-def main(command: Annotated[Optional[str], "command to run"] = typer.Argument(None)):
+def main(command: Optional[str] = typer.Argument(None)):
     match command:
         case "tfi":
             tfi()
@@ -150,8 +164,7 @@ def main(command: Annotated[Optional[str], "command to run"] = typer.Argument(No
         case "add-ip":
             add_ip()
         case "rm-ip":
-            ip = typer.prompt("IP address to remove")
-            rm_ip(ip)
+            rm_ip()
         case "hello":
             hello_world()
         case _: