From 9a40a5d0804fcbbefd1d41ce557bfce64aac2bc2 Mon Sep 17 00:00:00 2001 From: Dan Peterson Date: Thu, 14 Nov 2024 07:08:06 -0400 Subject: [PATCH] add pipe_once special command Mostly copied from mycli. Add tests for once as part of getting that going. Updates #192 --- litecli/main.py | 2 + litecli/packages/special/iocommands.py | 74 +++++++++++++++++++++----- tests/test_special_iocommands.py | 59 ++++++++++++++++++++ 3 files changed, 121 insertions(+), 14 deletions(-) create mode 100644 tests/test_special_iocommands.py diff --git a/litecli/main.py b/litecli/main.py index d1d3713..e3e6cc7 100644 --- a/litecli/main.py +++ b/litecli/main.py @@ -472,6 +472,7 @@ def one_iteration(text=None): result_count += 1 mutating = mutating or is_mutating(status) special.unset_once_if_written() + special.unset_pipe_once_if_written() except EOFError as e: raise e except KeyboardInterrupt: @@ -658,6 +659,7 @@ def output(self, output, status=None): self.log_output(line) special.write_tee(line) special.write_once(line) + special.write_pipe_once(line) if fits or output_via_pager: # buffering diff --git a/litecli/packages/special/iocommands.py b/litecli/packages/special/iocommands.py index 78f8707..eeba814 100644 --- a/litecli/packages/special/iocommands.py +++ b/litecli/packages/special/iocommands.py @@ -21,7 +21,10 @@ use_expanded_output = False PAGER_ENABLED = True tee_file = None -once_file = once_file_args = written_to_once_file = None +once_file = None +written_to_once_file = None +pipe_once_process = None +written_to_pipe_once_process = False favoritequeries = FavoriteQueries(ConfigObj()) @@ -376,9 +379,12 @@ def write_tee(output): aliases=("\\o", "\\once"), ) def set_once(arg, **_): - global once_file_args - - once_file_args = parseargfile(arg) + global once_file, written_to_once_file + try: + once_file = open(**parseargfile(arg)) + except (IOError, OSError) as e: + raise OSError("Cannot write to file '{}': {}".format(e.filename, e.strerror)) + written_to_once_file = False return [(None, None, None, "")] @@ -386,26 +392,66 @@ def set_once(arg, **_): @export def write_once(output): global once_file, written_to_once_file - if output and once_file_args: - if once_file is None: - try: - once_file = open(**once_file_args) - except (IOError, OSError) as e: - once_file = None - raise OSError("Cannot write to file '{}': {}".format(e.filename, e.strerror)) - + if output and once_file: click.echo(output, file=once_file, nl=False) click.echo("\n", file=once_file, nl=False) + once_file.flush() written_to_once_file = True @export def unset_once_if_written(): """Unset the once file, if it has been written to.""" - global once_file, once_file_args, written_to_once_file + global once_file, written_to_once_file if once_file and written_to_once_file: once_file.close() - once_file = once_file_args = written_to_once_file = None + once_file = written_to_once_file = None + + +@special_command("\\pipe_once", "\\| command", "Send next result to a subprocess.", aliases=("\\|",)) +def set_pipe_once(arg, **_): + global pipe_once_process, written_to_pipe_once_process + pipe_once_cmd = shlex.split(arg) + if len(pipe_once_cmd) == 0: + raise OSError("pipe_once requires a command") + written_to_pipe_once_process = False + pipe_once_process = subprocess.Popen( + pipe_once_cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + bufsize=1, + encoding="UTF-8", + universal_newlines=True, + ) + return [(None, None, None, "")] + + +@export +def write_pipe_once(output): + global pipe_once_process, written_to_pipe_once_process + if output and pipe_once_process: + try: + click.echo(output, file=pipe_once_process.stdin, nl=False) + click.echo("\n", file=pipe_once_process.stdin, nl=False) + except (IOError, OSError) as e: + pipe_once_process.terminate() + raise OSError("Failed writing to pipe_once subprocess: {}".format(e.strerror)) + written_to_pipe_once_process = True + + +@export +def unset_pipe_once_if_written(): + """Unset the pipe_once cmd, if it has been written to.""" + global pipe_once_process, written_to_pipe_once_process + if written_to_pipe_once_process: + (stdout_data, stderr_data) = pipe_once_process.communicate() + if len(stdout_data) > 0: + print(stdout_data.rstrip("\n")) + if len(stderr_data) > 0: + print(stderr_data.rstrip("\n")) + pipe_once_process = None + written_to_pipe_once_process = False @special_command( diff --git a/tests/test_special_iocommands.py b/tests/test_special_iocommands.py new file mode 100644 index 0000000..87ccc32 --- /dev/null +++ b/tests/test_special_iocommands.py @@ -0,0 +1,59 @@ +import os +import tempfile + +import pytest + +import litecli.packages.special + + +def test_once_command(): + with pytest.raises(TypeError): + litecli.packages.special.execute(None, ".once") + + with pytest.raises(OSError): + litecli.packages.special.execute(None, ".once /proc/access-denied") + + litecli.packages.special.write_once("hello world") # write without file set + # keep Windows from locking the file with delete=False + with tempfile.NamedTemporaryFile(delete=False) as f: + litecli.packages.special.execute(None, ".once " + f.name) + litecli.packages.special.write_once("hello world") + if os.name == "nt": + assert f.read() == b"hello world\r\n" + else: + assert f.read() == b"hello world\n" + + litecli.packages.special.execute(None, ".once -o " + f.name) + litecli.packages.special.write_once("hello world line 1") + litecli.packages.special.write_once("hello world line 2") + f.seek(0) + if os.name == "nt": + assert f.read() == b"hello world line 1\r\nhello world line 2\r\n" + else: + assert f.read() == b"hello world line 1\nhello world line 2\n" + # delete=False means we should try to clean up + try: + if os.path.exists(f.name): + os.remove(f.name) + except Exception as e: + print(f"An error occurred while attempting to delete the file: {e}") + + +def test_pipe_once_command(): + with pytest.raises(IOError): + litecli.packages.special.execute(None, "\\pipe_once") + + with pytest.raises(OSError): + litecli.packages.special.execute(None, "\\pipe_once /proc/access-denied") + + if os.name == "nt": + litecli.packages.special.execute(None, '\\pipe_once python -c "import sys; print(len(sys.stdin.read().strip()))"') + litecli.packages.special.write_pipe_once("hello world") + litecli.packages.special.unset_pipe_once_if_written() + else: + with tempfile.NamedTemporaryFile(delete=False) as f: + litecli.packages.special.execute(None, "\\pipe_once tee " + f.name) + litecli.packages.special.write_pipe_once("hello world") + litecli.packages.special.unset_pipe_once_if_written() + f.seek(0) + assert f.read() == b"hello world\n"