sh_run.py

import subprocess as sp
import textwrap

cmd = 'squeue --cluster=testbed --Format=jobid:10,partition:15,username:15,account:8,timeused:15,nodelist:15,name:50'
cmd = 'env | grep -i VENV'
cmd = 'python -c \'import sys; sys.stdout.write("stdout\\n"); sys.stderr.write("stderr\\n")\''
cmd = 'echo -e ""; sleep 1 && echo -e \'fal\\nse\' && false'

#sp.run(cmd, shell=1)

#proc = sp.Popen(cmd, shell=1, stdout=sp.PIPE, stderr=sp.PIPE)
#   proc = sp.Popen(cmd, shell=1, stdout=sp.PIPE, stderr=sp.STDOUT)
#   while line := proc.stdout.readline().decode():
#       print('[INFO]', repr(line))
#   
#   #print('[ERROR]', proc.stderr.readline().decode().strip())
#   print(proc.poll())
#   print('waiting')
#   returncode = proc.wait()
#   print(repr(proc.stdout.read().decode()))
#   print(returncode, proc.returncode)
#   exit()

def sh(cmd, *, capture_output=False, check=False, cwd=None, env=None):
    """
    If `capture_output`, run as `sp.run`.
    Otherwise, flush output (2>&1) and log lines to trace,
    where file handler is trace level, and console handler is info/debug level.
    """
    cmd = textwrap.dedent(cmd).strip()
    logger.debug('command:\n%s', cmd)

    # TODO: not completed yet
    if 1 or capture_output:
        print(f'{capture_output=}')
        return sp.run(cmd, shell=True, check=check, capture_output=capture_output, cwd=cwd, env=env)

    with sp.Popen(cmd, shell=True, stdout=sp.PIPE, stderr=sp.STDOUT) as proc:
        stdout_lines = []
        while line := proc.stdout.readline():
            logger.info(line)
            stdout_lines.append(line)

    stdout = b''.join(stdout_lines)
    if check and proc.returncode:
        raise sp.CalledProcessError(proc.returncode, proc.args, stdout)

    return sp.CompletedProcess(proc.args, proc.returncode, stdout)


class ShRun:
    """
    Run shell commands in subshell, print console, and log if given.
    """

    def __init__(self, logger=None):
        self.logger = logger

    def __call__(self, cmd, check=False):
        """
        Always capture stdout and redirect stderr to stdout.
        Log each line if logger is given.
        """
        cmd = textwrap.dedent(cmd)

        if self.logger:
            with sp.Popen(cmd, shell=True, stdout=sp.PIPE, stderr=sp.STDOUT) as proc:
                stdout_lines = []
                while line := proc.stdout.readline():
                    stdout_lines.append(line)
                    self.logger.info(line)
            completed_process = sp.CompletedProcess(proc.args, proc.returncode, b''.join(stdout_lines))
        else:
            completed_process = sp.run(cmd, shell=True, check=check, stdout=sp.PIPE, stderr=sp.STDOUT)

        return completed_process


proc = ShRun()(cmd)
print(f'{proc=}')
print(f'{proc.returncode=}')
print(f'{proc.stdout=}')
print(f'{proc.stderr=}')

print('==========')

import logging
logging.basicConfig(level=logging.INFO)
proc = ShRun(logging.getLogger())(cmd)
print(f'{proc=}')
print(f'{proc.returncode=}')
print(f'{proc.stdout=}')
print(f'{proc.stderr=}')

print('==========')

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
proc = sh(cmd)
print(f'{proc=}')
print(f'{proc.returncode=}')
print(f'{proc.stdout=}')
print(f'{proc.stderr=}')

print('==========')

proc = sh(cmd, capture_output=True)
print(f'{proc=}')
print(f'{proc.returncode=}')
print(f'{proc.stdout=}')
print(f'{proc.stderr=}')

print('==========1')

try:
    proc = sh(cmd, capture_output=True, check=True)
except Exception as e:
    print(e.output)

print('==========2')

try:
    proc = sh(cmd, check=True)
except Exception as e:
    print(e.output)

split_squeue

import subprocess as sp

cmd = 'squeue --cluster=testbed --Format=jobid:10,partition:20,username:15,account:8,timeused:15,nodelist:15,name:50'
proc = sp.run(cmd, shell=1, capture_output=1, check=1)
for line in proc.stdout.decode().splitlines():
    jobid = line[:10].strip()
    partition = line[10:30].strip()
    username = line[30:45].strip()
    account = line[45:53].strip()
    timeused = line[53:68].strip()
    nodelist = line[68:83].strip()
    name = line[83:133].strip()
    print(f'{jobid}, {partition}, {username}, {account}, {timeused}, {nodelist}, {name}')