109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
import argparse
|
|
import abc
|
|
from typing import List
|
|
from dataclasses import dataclass
|
|
from textwrap import dedent
|
|
from urllib import parse as urlparse
|
|
|
|
@dataclass
|
|
class PipeEnd:
|
|
name: str
|
|
format: str
|
|
|
|
@dataclass
|
|
class ShellMapping:
|
|
code: str
|
|
output_name: str
|
|
output_format: str
|
|
input_format: str
|
|
imports: List[str]
|
|
|
|
|
|
class AbstractMapper(abc.ABC):
|
|
@abc.abstractmethod
|
|
def __init__(self, command: List[str]):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def get_input_types(self):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def generate(self, input_name: str, output_formats: List[str]) -> ShellMapping:
|
|
pass
|
|
|
|
class CurlMapper(AbstractMapper):
|
|
def __init__(self, command):
|
|
parser = argparse.ArgumentParser("curl", exit_on_error=False)
|
|
parser.add_argument("url")
|
|
self._parsed = parser.parse_args(command[1:])
|
|
|
|
def get_input_types(self):
|
|
return []
|
|
|
|
def generate(self, input_name: str, output_formats: List[str]) -> ShellMapping:
|
|
url = urlparse.urlparse(self._parsed.url)
|
|
query = urlparse.parse_qs(url.query)
|
|
url_no_qs = urlparse.urlunparse(url._replace(query=None))
|
|
code = dedent(f"""\
|
|
params = {query!r}
|
|
res = requests.get({url_no_qs!r}, params=params).text
|
|
""")
|
|
return ShellMapping(code=code, output_name="res", output_format="str", input_format="", imports=["import requests"])
|
|
|
|
|
|
class ShellMapper(AbstractMapper):
|
|
def __init__(self, command):
|
|
self.command = command
|
|
|
|
def get_input_types(self):
|
|
return ["bytes"]
|
|
|
|
def generate(self, input_name: str, output_formats: List[str]) -> ShellMapping:
|
|
out_var = f"out_{self.command[0]}" # name var after command
|
|
if input_name:
|
|
code = dedent(f"""\
|
|
proc = subprocess.Popen({self.command!r}, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=sys.stderr)
|
|
(stdout, _) = proc.communicate({input_name}.encode("utf-8"))
|
|
{out_var} = stdout.decode("utf-8")
|
|
""")
|
|
else:
|
|
code = f"""{out_var} = subprocess.run({self.command!r}, check=True, stdout=subprocess.PIPE, stderr=sys.stderr).stdout"""
|
|
|
|
return ShellMapping(
|
|
code=code,
|
|
imports=["import subprocess", "import sys"],
|
|
input_format="bytes",
|
|
output_format="bytes",
|
|
output_name=out_var,
|
|
)
|
|
|
|
|
|
class XargsMapper(AbstractMapper):
|
|
def __init__(self, command):
|
|
parser = argparse.ArgumentParser("xargs", exit_on_error=False)
|
|
parser.add_argument("command", nargs="+")
|
|
parser.add_argument("-L", type=int)
|
|
self._parsed = parser.parse_args(command[1:])
|
|
self.input_types = ["list", "bytes"]
|
|
|
|
def get_input_types(self):
|
|
return ["list", "bytes"]
|
|
|
|
def generate(self, input_name: str, output_formats: List[str]) -> ShellMapping:
|
|
imports = []
|
|
if self._parsed.L >= 1:
|
|
command_str = self._parsed.command
|
|
run_command = "print(i) # TODO"
|
|
if True:
|
|
input_name = f"shlex.split({input_name}, posix=True)"
|
|
imports.append("import shlex")
|
|
code = dedent(
|
|
f"""\
|
|
for i in {input_name}:
|
|
{run_command}"""
|
|
)
|
|
else:
|
|
code = "# todo: shell stuff"
|
|
return ShellMapping(code=code, output_name="", output_format="", input_format="list", imports=imports)
|