1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
|
from __future__ import annotations
import os
import sys
import shutil
import tarfile
import platform
import subprocess
from typing import TYPE_CHECKING, List
from pathlib import Path
from argparse import ArgumentParser
import httpx
from .._errors import CLIError, SilentCLIError
from .._models import BaseModel
if TYPE_CHECKING:
from argparse import _SubParsersAction
def register(subparser: _SubParsersAction[ArgumentParser]) -> None:
sub = subparser.add_parser("migrate")
sub.set_defaults(func=migrate, args_model=MigrateArgs, allow_unknown_args=True)
sub = subparser.add_parser("grit")
sub.set_defaults(func=grit, args_model=GritArgs, allow_unknown_args=True)
class GritArgs(BaseModel):
# internal
unknown_args: List[str] = []
def grit(args: GritArgs) -> None:
grit_path = install()
try:
subprocess.check_call([grit_path, *args.unknown_args])
except subprocess.CalledProcessError:
# stdout and stderr are forwarded by subprocess so an error will already
# have been displayed
raise SilentCLIError() from None
class MigrateArgs(BaseModel):
# internal
unknown_args: List[str] = []
def migrate(args: MigrateArgs) -> None:
grit_path = install()
try:
subprocess.check_call([grit_path, "apply", "openai", *args.unknown_args])
except subprocess.CalledProcessError:
# stdout and stderr are forwarded by subprocess so an error will already
# have been displayed
raise SilentCLIError() from None
# handles downloading the Grit CLI until they provide their own PyPi package
KEYGEN_ACCOUNT = "custodian-dev"
def _cache_dir() -> Path:
xdg = os.environ.get("XDG_CACHE_HOME")
if xdg is not None:
return Path(xdg)
return Path.home() / ".cache"
def _debug(message: str) -> None:
if not os.environ.get("DEBUG"):
return
sys.stdout.write(f"[DEBUG]: {message}\n")
def install() -> Path:
"""Installs the Grit CLI and returns the location of the binary"""
if sys.platform == "win32":
raise CLIError("Windows is not supported yet in the migration CLI")
_debug("Using Grit installer from GitHub")
platform = "apple-darwin" if sys.platform == "darwin" else "unknown-linux-gnu"
dir_name = _cache_dir() / "openai-python"
install_dir = dir_name / ".install"
target_dir = install_dir / "bin"
target_path = target_dir / "grit"
temp_file = target_dir / "grit.tmp"
if target_path.exists():
_debug(f"{target_path} already exists")
sys.stdout.flush()
return target_path
_debug(f"Using Grit CLI path: {target_path}")
target_dir.mkdir(parents=True, exist_ok=True)
if temp_file.exists():
temp_file.unlink()
arch = _get_arch()
_debug(f"Using architecture {arch}")
file_name = f"grit-{arch}-{platform}"
download_url = f"https://github.com/getgrit/gritql/releases/latest/download/{file_name}.tar.gz"
sys.stdout.write(f"Downloading Grit CLI from {download_url}\n")
with httpx.Client() as client:
download_response = client.get(download_url, follow_redirects=True)
if download_response.status_code != 200:
raise CLIError(f"Failed to download Grit CLI from {download_url}")
with open(temp_file, "wb") as file:
for chunk in download_response.iter_bytes():
file.write(chunk)
unpacked_dir = target_dir / "cli-bin"
unpacked_dir.mkdir(parents=True, exist_ok=True)
with tarfile.open(temp_file, "r:gz") as archive:
if sys.version_info >= (3, 12):
archive.extractall(unpacked_dir, filter="data")
else:
archive.extractall(unpacked_dir)
_move_files_recursively(unpacked_dir, target_dir)
shutil.rmtree(unpacked_dir)
os.remove(temp_file)
os.chmod(target_path, 0o755)
sys.stdout.flush()
return target_path
def _move_files_recursively(source_dir: Path, target_dir: Path) -> None:
for item in source_dir.iterdir():
if item.is_file():
item.rename(target_dir / item.name)
elif item.is_dir():
_move_files_recursively(item, target_dir)
def _get_arch() -> str:
architecture = platform.machine().lower()
# Map the architecture names to Grit equivalents
arch_map = {
"x86_64": "x86_64",
"amd64": "x86_64",
"armv7l": "aarch64",
"arm64": "aarch64",
}
return arch_map.get(architecture, architecture)
|
