Merge pull request #24 from gormaniac/utils-refactor

refactor: move more functions around

src/stormlibpp/import.py

@@ -3,11 +3,8 @@
 
 import argparse
 import asyncio
-import contextlib
 import copy
-import csv
 import functools
-import json
 import os
 import pathlib
 import re
@@ -22,38 +19,8 @@
 from .httpcore import HttpCortex
 from .output import handle_msg, log_storm_msg, OUTP
 from .stormcli import start_storm_cli
-from .utils import get_cortex_creds
-
-
-def endswith(text, items):
-    for item in items:
-        if text.endswith(item):
-            return True
-    return False
-
-
-def csv_genr(fd):
-    for row in csv.reader(fd):
-        yield row
-
-
-def json_genr(fd):
-    # TODO - Support nested key to start with
-    data = json.load(fd)
-
-    if type(data, dict):
-        for key, row in data.items():
-            yield [key, row]
-    else:
-        count = 0
-        for row in data:
-            yield [count, row]
-            count += 1
-
-
-def txt_genr(fd):
-    for row in fd.readlines():
-        yield row.strip()
+from .telepath import tpath_proxy_contextmanager
+from .utils import csv_genr, endswith, json_genr, get_cortex_creds, txt_genr
 
 
 async def import_storm_text(core, text, stormopts, print_skips, logfd):
@@ -135,7 +102,7 @@ async def add_data(self):
         return nodecount
 
 
-def find_data_files(orig: str, files: list[str]):
+def find_data_files(orig: str, files: list[str]) -> tuple[str, str]:
     """
     Find all data files in a list of files that have the same name (and path) as the original file.
     """
@@ -164,13 +131,6 @@ def find_data_files(orig: str, files: list[str]):
     return data_files
 
 
-# TODO - Keep this in this file?
-# Maybe put in teletpath.py and rename to tpath_proxy_contextmanager
-@contextlib.contextmanager
-def cortex_proxy_contextmanager(cortex_url):
-    yield s_telepath.openurl(cortex_url)
-
-
 def get_args(argv: list[str]):
     parser = argparse.ArgumentParser(description=__doc__, parents=[USER_PARSER])
     parser.add_argument(
@@ -256,7 +216,7 @@ async def get_cortex_obj(cortex, local, http, no_verify, user):
                 HttpCortex, cortex, username, password, ssl_verify=not no_verify
             )
         else:
-            core_obj = functools.partial(cortex_proxy_contextmanager, cortex)
+            core_obj = functools.partial(tpath_proxy_contextmanager, cortex)
     elif local:
         core_obj = s_cortex.getTempCortex
     else:

src/stormlibpp/telepath.py

@@ -1,8 +1,11 @@
 """Python objects for working with Synapse Telepath endpoints."""
 
 
+import contextlib
 from typing import Any, TypedDict
 
+import synapse.telepath as s_telepath
+
 
 class TelepathRetn(TypedDict):
     """A standard return value from a Telepath API endpoint.
@@ -60,3 +63,10 @@ def genDefaultTelepathRetn(obj=TelepathRetn, /, default_data=None) -> TelepathRe
     """
 
     return obj(status=True, mesg="", data=default_data)
+
+
+@contextlib.contextmanager
+def tpath_proxy_contextmanager(cortex_url: str) -> s_telepath.Proxy:
+    """Context manager that yields a telepath proxy object connected to cortex_url."""
+
+    yield s_telepath.openurl(cortex_url)

src/stormlibpp/utils.py

@@ -1,7 +1,9 @@
 """Miscellaneous functions that are helpful for working with Storm/Synapse."""
 
 
+import csv
 import getpass
+import json
 import os
 
 import synapse.exc as s_exc
@@ -44,6 +46,67 @@ def chk_storm_syntax(storm: str) -> None:
         raise errors.StormSyntaxError("Storm Syntax Error!", err) from err
 
 
+def csv_genr(fd: str) -> list:
+    """Read a CSV from fd and yield each row."""
+
+    for row in csv.reader(fd):
+        yield row
+
+
+def endswith(text: str, items: list[str]) -> bool:
+    """Check if text ends with any of items."""
+
+    for item in items:
+        if text.endswith(item):
+            return True
+    return False
+
+
+def get_cortex_creds(_user: str | None = None) -> tuple[str, str]:
+    """Get credentials to use when connecting to a Cortex over HTTP.
+    
+    If ``_user`` is set, it overrides any other options to get a username.
+    Otherwise use the value of the ``CORTEX_USER`` environment variable. If that
+    is also empty, prompt for a username. A default option is given to the user
+    at the prompt, it is the return of ``getpass.getuser()``.
+
+    The password is read from the ``CORTEX_PASS`` environment variable and the
+    user is prompted if this is empty.
+    """
+
+    if _user:
+        username = _user
+    elif envusr := os.environ.get("CORTEX_USER"):
+        username = envusr
+    else:
+        gpusr = getpass.getuser()
+        inusr = input(f"Username [{gpusr}]: ")
+        username = inusr if inusr else gpusr
+
+    if envpw := os.environ.get("CORTEX_PASS"):
+        password = envpw
+    else:
+        password = getpass.getpass()
+
+    return username, password
+
+
+def json_genr(fd):
+    """Read a JSON file from fd, and yield each top-level item within it."""
+
+    # TODO - Support nested key to start with
+    data = json.load(fd)
+
+    if type(data, dict):
+        for key, row in data.items():
+            yield [key, row]
+    else:
+        count = 0
+        for row in data:
+            yield [count, row]
+            count += 1
+
+
 def normver(ver: str | tuple) -> tuple[str, tuple]:
     """Take either a version str "x.x.x" or tuple (x, x, x) and return both.
 
@@ -65,19 +128,8 @@ def normver(ver: str | tuple) -> tuple[str, tuple]:
     return (verstr, vertup)
 
 
-def get_cortex_creds(_user):
-    if _user:
-        username = _user
-    elif envusr := os.environ.get("CORTEX_USER"):
-        username = envusr
-    else:
-        gpusr = getpass.getuser()
-        inusr = input(f"Username [{gpusr}]: ")
-        username = inusr if inusr else gpusr
-
-    if envpw := os.environ.get("CORTEX_PASS"):
-        password = envpw
-    else:
-        password = getpass.getpass()
+def txt_genr(fd):
+    """Read a raw text file from fd and yield each line."""
 
-    return username, password
+    for row in fd.readlines():
+        yield row.strip()