Merge pull request #1181 from datajoint/cli-1095-continued

DataJoint CLI continued (#1095)

Author: dimitri-yatsenko

CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 0.14.3 -- TBD
 - Added - `dj.Top` restriction ([#1024](https://github.com/datajoint/datajoint-python/issues/1024)) PR [#1084](https://github.com/datajoint/datajoint-python/pull/1084)
 - Fixed - Added encapsulating double quotes to comply with [DOT language](https://graphviz.org/doc/info/lang.html) - PR [#1177](https://github.com/datajoint/datajoint-python/pull/1177)
+- Added - Datajoint python CLI ([#940](https://github.com/datajoint/datajoint-python/issues/940)) PR [#1095](https://github.com/datajoint/datajoint-python/pull/1095)
 - Added - Ability to set hidden attributes on a table - PR [#1091](https://github.com/datajoint/datajoint-python/pull/1091)
 
 ### 0.14.2 -- Aug 19, 2024

datajoint/__init__.py

@@ -52,6 +52,7 @@
     "key",
     "key_hash",
     "logger",
+    "cli",
 ]
 
 from .logging import logger
@@ -71,6 +72,7 @@
 from .attribute_adapter import AttributeAdapter
 from . import errors
 from .errors import DataJointError
+from .cli import cli
 
 ERD = Di = Diagram  # Aliases for Diagram
 schema = Schema  # Aliases for Schema

datajoint/cli.py

@@ -0,0 +1,77 @@
+import argparse
+from code import interact
+from collections import ChainMap
+import datajoint as dj
+
+
+def cli(args: list = None):
+    """
+    Console interface for DataJoint Python
+
+    :param args: List of arguments to be passed in, defaults to reading stdin
+    :type args: list, optional
+    """
+    parser = argparse.ArgumentParser(
+        prog="datajoint",
+        description="DataJoint console interface.",
+        conflict_handler="resolve",
+    )
+    parser.add_argument(
+        "-V", "--version", action="version", version=f"{dj.__name__} {dj.__version__}"
+    )
+    parser.add_argument(
+        "-u",
+        "--user",
+        type=str,
+        default=dj.config["database.user"],
+        required=False,
+        help="Datajoint username",
+    )
+    parser.add_argument(
+        "-p",
+        "--password",
+        type=str,
+        default=dj.config["database.password"],
+        required=False,
+        help="Datajoint password",
+    )
+    parser.add_argument(
+        "-h",
+        "--host",
+        type=str,
+        default=dj.config["database.host"],
+        required=False,
+        help="Datajoint host",
+    )
+    parser.add_argument(
+        "-s",
+        "--schemas",
+        nargs="+",
+        type=str,
+        required=False,
+        help="A list of virtual module mappings in `db:schema ...` format",
+    )
+    kwargs = vars(parser.parse_args(args))
+    mods = {}
+    if kwargs["user"]:
+        dj.config["database.user"] = kwargs["user"]
+    if kwargs["password"]:
+        dj.config["database.password"] = kwargs["password"]
+    if kwargs["host"]:
+        dj.config["database.host"] = kwargs["host"]
+    if kwargs["schemas"]:
+        for vm in kwargs["schemas"]:
+            d, m = vm.split(":")
+            mods[m] = dj.create_virtual_module(m, d)
+
+    banner = "dj repl\n"
+    if mods:
+        modstr = "\n".join("  - {}".format(m) for m in mods)
+        banner += "\nschema modules:\n\n" + modstr + "\n"
+    interact(banner, local=dict(ChainMap(mods, locals(), globals())))
+
+    raise SystemExit
+
+
+if __name__ == "__main__":
+    cli()

setup.py

@@ -39,6 +39,9 @@
         "automated research workflows",
     ],
     packages=find_packages(exclude=["contrib", "docs", "tests*"]),
+    entry_points={
+        "console_scripts": ["dj=datajoint.cli:cli", "datajoint=datajoint.cli:cli"],
+    },
     install_requires=requirements,
     python_requires="~={}.{}".format(*min_py_version),
     setup_requires=["otumat"],  # maybe remove due to conflicts?

tests/test_cli.py

@@ -0,0 +1,116 @@
+"""
+Collection of test cases to test the dj cli
+"""
+
+import json
+import ast
+import subprocess
+import pytest
+import datajoint as dj
+
+
+def test_cli_version(capsys):
+    with pytest.raises(SystemExit) as pytest_wrapped_e:
+        dj.cli(args=["-V"])
+    assert pytest_wrapped_e.type == SystemExit
+    assert pytest_wrapped_e.value.code == 0
+
+    captured_output = capsys.readouterr().out
+    assert captured_output == f"{dj.__name__} {dj.__version__}\n"
+
+
+def test_cli_help(capsys):
+    with pytest.raises(SystemExit) as pytest_wrapped_e:
+        dj.cli(args=["--help"])
+    assert pytest_wrapped_e.type == SystemExit
+    assert pytest_wrapped_e.value.code == 0
+
+    captured_output = capsys.readouterr().out
+    assert captured_output.strip()
+
+
+def test_cli_config():
+    process = subprocess.Popen(
+        ["dj"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+
+    process.stdin.write("dj.config\n")
+    process.stdin.flush()
+
+    stdout, stderr = process.communicate()
+    cleaned = stdout.strip(" >\t\n\r")
+    for key in ("database.user", "database.password", "database.host"):
+        assert key in cleaned, f"Key {key} not found in config from stdout: {cleaned}"
+
+
+def test_cli_args():
+    process = subprocess.Popen(
+        ["dj", "-utest_user", "-ptest_pass", "-htest_host"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+
+    process.stdin.write("dj.config['database.user']\n")
+    process.stdin.write("dj.config['database.password']\n")
+    process.stdin.write("dj.config['database.host']\n")
+    process.stdin.flush()
+
+    stdout, stderr = process.communicate()
+    assert "test_user" == stdout[5:14]
+    assert "test_pass" == stdout[21:30]
+    assert "test_host" == stdout[37:46]
+
+
+def test_cli_schemas(prefix, connection_root):
+    schema = dj.Schema(prefix + "_cli", locals(), connection=connection_root)
+
+    @schema
+    class IJ(dj.Lookup):
+        definition = """  # tests restrictions
+        i  : int
+        j  : int
+        """
+        contents = list(dict(i=i, j=j + 2) for i in range(3) for j in range(3))
+
+    process = subprocess.Popen(
+        ["dj", "-s", "djtest_cli:test_schema"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+
+    process.stdin.write("test_schema.__dict__['__name__']\n")
+    process.stdin.write("test_schema.__dict__['schema']\n")
+    process.stdin.write("test_schema.IJ.fetch(as_dict=True)\n")
+    process.stdin.flush()
+
+    stdout, stderr = process.communicate()
+    fetch_res = [
+        {"i": 0, "j": 2},
+        {"i": 0, "j": 3},
+        {"i": 0, "j": 4},
+        {"i": 1, "j": 2},
+        {"i": 1, "j": 3},
+        {"i": 1, "j": 4},
+        {"i": 2, "j": 2},
+        {"i": 2, "j": 3},
+        {"i": 2, "j": 4},
+    ]
+    assert (
+        "\
+dj repl\n\n\
+\
+schema modules:\n\n\
+  - test_schema"
+        == stderr[159:200]
+    )
+    assert "'test_schema'" == stdout[4:17]
+    assert "Schema `djtest_cli`" == stdout[22:41]
+    assert fetch_res == json.loads(stdout[47:209].replace("'", '"'))