kscale 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
kscale/__init__.py CHANGED
@@ -1 +1 @@
1
- __version__ = "0.0.2"
1
+ __version__ = "0.0.3"
kscale/store/cli.py CHANGED
@@ -1,12 +1,13 @@
1
1
  """Defines the top-level KOL CLI."""
2
2
 
3
3
  import argparse
4
+ import asyncio
4
5
  from typing import Sequence
5
6
 
6
7
  from kscale.store import pybullet, urdf
7
8
 
8
9
 
9
- def main(args: Sequence[str] | None = None) -> None:
10
+ async def main(args: Sequence[str] | None = None) -> None:
10
11
  parser = argparse.ArgumentParser(description="K-Scale OnShape Library", add_help=False)
11
12
  parser.add_argument(
12
13
  "subcommand",
@@ -20,11 +21,15 @@ def main(args: Sequence[str] | None = None) -> None:
20
21
 
21
22
  match parsed_args.subcommand:
22
23
  case "urdf":
23
- urdf.main(remaining_args)
24
+ await urdf.main(remaining_args)
24
25
  case "pybullet":
25
- pybullet.main(remaining_args)
26
+ await pybullet.main(remaining_args)
27
+
28
+
29
+ def sync_main(args: Sequence[str] | None = None) -> None:
30
+ asyncio.run(main(args))
26
31
 
27
32
 
28
33
  if __name__ == "__main__":
29
34
  # python3 -m kscale.store.cli
30
- main()
35
+ sync_main()
kscale/store/pybullet.py CHANGED
@@ -1,6 +1,7 @@
1
1
  """Simple script to interact with a URDF in PyBullet."""
2
2
 
3
3
  import argparse
4
+ import asyncio
4
5
  import itertools
5
6
  import logging
6
7
  import math
@@ -8,12 +9,14 @@ import time
8
9
  from pathlib import Path
9
10
  from typing import Sequence
10
11
 
12
+ from kscale.store.urdf import download_urdf
13
+
11
14
  logger = logging.getLogger(__name__)
12
15
 
13
16
 
14
- def main(args: Sequence[str] | None = None) -> None:
17
+ async def main(args: Sequence[str] | None = None) -> None:
15
18
  parser = argparse.ArgumentParser(description="Show a URDF in PyBullet")
16
- parser.add_argument("urdf", nargs="?", help="Path to the URDF file")
19
+ parser.add_argument("listing_id", help="Listing ID for the URDF")
17
20
  parser.add_argument("--dt", type=float, default=0.01, help="Time step")
18
21
  parser.add_argument("-n", "--hide-gui", action="store_true", help="Hide the GUI")
19
22
  parser.add_argument("--no-merge", action="store_true", help="Do not merge fixed links")
@@ -23,6 +26,11 @@ def main(args: Sequence[str] | None = None) -> None:
23
26
  parser.add_argument("--show-collision", action="store_true", help="Show collision meshes")
24
27
  parsed_args = parser.parse_args(args)
25
28
 
29
+ # Gets the URDF path.
30
+ urdf_path = await download_urdf(parsed_args.listing_id)
31
+
32
+ breakpoint()
33
+
26
34
  try:
27
35
  import pybullet as p # type: ignore[import-not-found]
28
36
  except ImportError:
@@ -46,13 +54,6 @@ def main(args: Sequence[str] | None = None) -> None:
46
54
  # Loads the floor plane.
47
55
  floor = p.loadURDF(str((Path(__file__).parent / "bullet" / "plane.urdf").resolve()))
48
56
 
49
- urdf_path = Path("robot" if parsed_args.urdf is None else parsed_args.urdf)
50
- if urdf_path.is_dir():
51
- try:
52
- urdf_path = next(urdf_path.glob("*.urdf"))
53
- except StopIteration:
54
- raise FileNotFoundError(f"No URDF files found in {urdf_path}")
55
-
56
57
  # Load the robot URDF.
57
58
  start_position = [0.0, 0.0, 1.0]
58
59
  start_orientation = p.getQuaternionFromEuler([0.0, 0.0, 0.0])
@@ -175,4 +176,4 @@ def main(args: Sequence[str] | None = None) -> None:
175
176
 
176
177
  if __name__ == "__main__":
177
178
  # python -m kscale.store.pybullet
178
- main()
179
+ asyncio.run(main())
kscale/store/urdf.py CHANGED
@@ -4,10 +4,11 @@ import argparse
4
4
  import asyncio
5
5
  import logging
6
6
  import os
7
+ import shutil
7
8
  import sys
8
9
  import tarfile
9
10
  from pathlib import Path
10
- from typing import Literal, Sequence
11
+ from typing import Literal, Sequence, get_args
11
12
 
12
13
  import httpx
13
14
  import requests
@@ -24,7 +25,7 @@ def get_api_key() -> str:
24
25
  api_key = Settings.load().store.api_key
25
26
  if not api_key:
26
27
  raise ValueError(
27
- "API key not found! Get one here and set it as the `KSCALE_API_KEY` environment variable or in your"
28
+ "API key not found! Get one here and set it as the `KSCALE_API_KEY` environment variable or in your "
28
29
  "config file: https://kscale.store/keys"
29
30
  )
30
31
  return api_key
@@ -34,6 +35,11 @@ def get_cache_dir() -> Path:
34
35
  return Path(Settings.load().store.cache_dir).expanduser().resolve()
35
36
 
36
37
 
38
+ def get_listing_dir(listing_id: str) -> Path:
39
+ (cache_dir := get_cache_dir() / listing_id).mkdir(parents=True, exist_ok=True)
40
+ return cache_dir
41
+
42
+
37
43
  def fetch_urdf_info(listing_id: str) -> UrdfResponse:
38
44
  url = f"https://api.kscale.store/urdf/info/{listing_id}"
39
45
  headers = {
@@ -44,14 +50,14 @@ def fetch_urdf_info(listing_id: str) -> UrdfResponse:
44
50
  return UrdfResponse(**response.json())
45
51
 
46
52
 
47
- async def download_artifact(artifact_url: str, cache_dir: Path) -> str:
53
+ async def download_artifact(artifact_url: str, cache_dir: Path) -> Path:
48
54
  filename = os.path.join(cache_dir, artifact_url.split("/")[-1])
49
55
  headers = {
50
56
  "Authorization": f"Bearer {get_api_key()}",
51
57
  }
52
58
 
53
59
  if not os.path.exists(filename):
54
- logger.info("Downloading artifact from %s" % artifact_url)
60
+ logger.info("Downloading artifact from %s", artifact_url)
55
61
 
56
62
  async with httpx.AsyncClient() as client:
57
63
  response = await client.get(artifact_url, headers=headers)
@@ -59,19 +65,18 @@ async def download_artifact(artifact_url: str, cache_dir: Path) -> str:
59
65
  with open(filename, "wb") as f:
60
66
  for chunk in response.iter_bytes(chunk_size=8192):
61
67
  f.write(chunk)
62
- logger.info("Artifact downloaded to %s" % filename)
68
+ logger.info("Artifact downloaded to %s", filename)
63
69
  else:
64
- logger.info("Artifact already cached at %s" % filename)
70
+ logger.info("Artifact already cached at %s", filename)
65
71
 
66
72
  # Extract the .tgz file
67
- extract_dir = os.path.join(cache_dir, os.path.splitext(os.path.basename(filename))[0])
68
- if not os.path.exists(extract_dir):
69
- logger.info(f"Extracting {filename} to {extract_dir}")
73
+ extract_dir = cache_dir / os.path.splitext(os.path.basename(filename))[0]
74
+ if not extract_dir.exists():
75
+ logger.info("Extracting %s to %s", filename, extract_dir)
70
76
  with tarfile.open(filename, "r:gz") as tar:
71
77
  tar.extractall(path=extract_dir)
72
- logger.info("Extraction complete")
73
78
  else:
74
- logger.info("Artifact already extracted at %s" % extract_dir)
79
+ logger.info("Artifact already extracted at %s", extract_dir)
75
80
 
76
81
  return extract_dir
77
82
 
@@ -84,8 +89,8 @@ def create_tarball(folder_path: str | Path, output_filename: str, cache_dir: Pat
84
89
  file_path = os.path.join(root, file)
85
90
  arcname = os.path.relpath(file_path, start=folder_path)
86
91
  tar.add(file_path, arcname=arcname)
87
- logger.info("Added %s as %s" % (file_path, arcname))
88
- logger.info("Created tarball %s" % tarball_path)
92
+ logger.info("Added %s as %s", file_path, arcname)
93
+ logger.info("Created tarball %s", tarball_path)
89
94
  return tarball_path
90
95
 
91
96
 
@@ -102,68 +107,102 @@ async def upload_artifact(tarball_path: str, listing_id: str) -> None:
102
107
 
103
108
  response.raise_for_status()
104
109
 
105
- logger.info("Uploaded artifact to %s" % url)
110
+ logger.info("Uploaded artifact to %s", url)
111
+
112
+
113
+ async def download_urdf(listing_id: str) -> Path:
114
+ try:
115
+ urdf_info = fetch_urdf_info(listing_id)
116
+
117
+ if urdf_info.urdf is None:
118
+ breakpoint()
119
+ raise ValueError(f"No URDF found for listing {listing_id}")
120
+
121
+ artifact_url = urdf_info.urdf.url
122
+ return await download_artifact(artifact_url, get_listing_dir(listing_id))
123
+
124
+ except requests.RequestException:
125
+ logger.exception("Failed to fetch URDF info")
126
+ raise
127
+
128
+
129
+ async def show_urdf_info(listing_id: str) -> None:
130
+ try:
131
+ urdf_info = fetch_urdf_info(listing_id)
132
+
133
+ if urdf_info.urdf:
134
+ logger.info("URDF Artifact ID: %s", urdf_info.urdf.artifact_id)
135
+ logger.info("URDF URL: %s", urdf_info.urdf.url)
136
+ else:
137
+ logger.info("No URDF found for listing %s", listing_id)
138
+ except requests.RequestException:
139
+ logger.exception("Failed to fetch URDF info")
140
+ raise
141
+
106
142
 
143
+ async def upload_urdf(listing_id: str, args: Sequence[str] | None = None) -> None:
144
+ parser = argparse.ArgumentParser(description="Upload a URDF artifact to the K-Scale store")
145
+ parser.add_argument("folder_path", help="The path to the folder containing the URDF files")
146
+ parsed_args = parser.parse_args(args)
147
+ folder_path = Path(parsed_args.folder_path).expanduser().resolve()
107
148
 
108
- def main(args: Sequence[str] | None = None) -> None:
149
+ output_filename = f"{listing_id}.tgz"
150
+ tarball_path = create_tarball(folder_path, output_filename, get_listing_dir(listing_id))
151
+
152
+ try:
153
+ fetch_urdf_info(listing_id)
154
+ await upload_artifact(tarball_path, listing_id)
155
+ except requests.RequestException:
156
+ logger.exception("Failed to upload artifact")
157
+ raise
158
+
159
+
160
+ async def remove_local_urdf(listing_id: str) -> None:
161
+ try:
162
+ if listing_id.lower() == "all":
163
+ cache_dir = get_cache_dir()
164
+ if cache_dir.exists():
165
+ logger.info("Removing all local caches at %s", cache_dir)
166
+ shutil.rmtree(cache_dir)
167
+ else:
168
+ logger.error("No local caches found")
169
+ else:
170
+ listing_dir = get_listing_dir(listing_id)
171
+ if listing_dir.exists():
172
+ logger.info("Removing local cache at %s", listing_dir)
173
+ shutil.rmtree(listing_dir)
174
+ else:
175
+ logger.error("No local cache found for listing %s", listing_id)
176
+
177
+ except Exception:
178
+ logger.error("Failed to remove local cache")
179
+ raise
180
+
181
+
182
+ Command = Literal["download", "info", "upload", "remove-local"]
183
+
184
+
185
+ async def main(args: Sequence[str] | None = None) -> None:
109
186
  parser = argparse.ArgumentParser(description="K-Scale URDF Store", add_help=False)
110
- parser.add_argument(
111
- "command",
112
- choices=["get", "info", "upload"],
113
- help="The command to run",
114
- )
187
+ parser.add_argument("command", choices=get_args(Command), help="The command to run")
115
188
  parser.add_argument("listing_id", help="The listing ID to operate on")
116
189
  parsed_args, remaining_args = parser.parse_known_args(args)
117
190
 
118
- command: Literal["get", "info", "upload"] = parsed_args.command
191
+ command: Command = parsed_args.command
119
192
  listing_id: str = parsed_args.listing_id
120
193
 
121
- def get_listing_dir() -> Path:
122
- (cache_dir := get_cache_dir() / listing_id).mkdir(parents=True, exist_ok=True)
123
- return cache_dir
124
-
125
194
  match command:
126
- case "get":
127
- try:
128
- urdf_info = fetch_urdf_info(listing_id)
129
-
130
- if urdf_info.urdf:
131
- artifact_url = urdf_info.urdf.url
132
- asyncio.run(download_artifact(artifact_url, get_listing_dir()))
133
- else:
134
- logger.info("No URDF found for listing %s" % listing_id)
135
- except requests.RequestException as e:
136
- logger.error("Failed to fetch URDF info: %s" % e)
137
- sys.exit(1)
195
+ case "download":
196
+ await download_urdf(listing_id)
138
197
 
139
198
  case "info":
140
- try:
141
- urdf_info = fetch_urdf_info(listing_id)
142
-
143
- if urdf_info.urdf:
144
- logger.info("URDF Artifact ID: %s" % urdf_info.urdf.artifact_id)
145
- logger.info("URDF URL: %s" % urdf_info.urdf.url)
146
- else:
147
- logger.info("No URDF found for listing %s" % listing_id)
148
- except requests.RequestException as e:
149
- logger.error("Failed to fetch URDF info: %s" % e)
150
- sys.exit(1)
199
+ await show_urdf_info(listing_id)
151
200
 
152
201
  case "upload":
153
- parser = argparse.ArgumentParser(description="Upload a URDF artifact to the K-Scale store")
154
- parser.add_argument("folder_path", help="The path to the folder containing the URDF files")
155
- parsed_args = parser.parse_args(remaining_args)
156
- folder_path = Path(parsed_args.folder_path).expanduser().resolve()
157
-
158
- output_filename = f"{listing_id}.tgz"
159
- tarball_path = create_tarball(folder_path, output_filename, get_listing_dir())
202
+ await upload_urdf(listing_id, remaining_args)
160
203
 
161
- try:
162
- urdf_info = fetch_urdf_info(listing_id)
163
- asyncio.run(upload_artifact(tarball_path, listing_id))
164
- except requests.RequestException as e:
165
- logger.error("Failed to upload artifact: %s" % e)
166
- sys.exit(1)
204
+ case "remove-local":
205
+ await remove_local_urdf(listing_id)
167
206
 
168
207
  case _:
169
208
  logger.error("Invalid command")
@@ -171,4 +210,4 @@ def main(args: Sequence[str] | None = None) -> None:
171
210
 
172
211
 
173
212
  if __name__ == "__main__":
174
- main()
213
+ asyncio.run(main())
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: kscale
3
- Version: 0.0.2
3
+ Version: 0.0.3
4
4
  Summary: The kscale project
5
5
  Home-page: https://github.com/kscalelabs/kscale
6
6
  Author: Benjamin Bolte
@@ -0,0 +1,15 @@
1
+ kscale/__init__.py,sha256=4GZKi13lDTD25YBkGakhZyEQZWTER_OWQMNPoH_UM2c,22
2
+ kscale/conf.py,sha256=dnO8qii7JeMPMdjfuxnFXURRM4krwzL404RnwDkyL2M,1441
3
+ kscale/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
4
+ kscale/store/__init__.py,sha256=sXLh7g3KC4QCFxcZGBTpG2scR7hmmBsMjq6LqRptkRg,22
5
+ kscale/store/cli.py,sha256=8ygg_1tZzOOHJotEIgSN9pfumcriPmA31sI_FCFQiTo,859
6
+ kscale/store/pybullet.py,sha256=y51Oc6ZjOGe38O-L-Lm6HDbgNWc-5mpeXCeXRtCUl18,7522
7
+ kscale/store/urdf.py,sha256=0u6c2UxYlKwgUl0n3VmMSmfTqdixNP5b-hjox5mIgqM,7064
8
+ kscale/store/gen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
9
+ kscale/store/gen/api.py,sha256=7r6KvlZyB-kWKSCixc-YOtaWPy2mn0v11d2tuLlRMSc,7700
10
+ kscale-0.0.3.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
11
+ kscale-0.0.3.dist-info/METADATA,sha256=KN2mHTYV9mkJFxYqe4u1PFkCFQkc5o-E3460k66gN_M,2028
12
+ kscale-0.0.3.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
13
+ kscale-0.0.3.dist-info/entry_points.txt,sha256=PaVs1ivqB0BBdGUsiFkxGUYjGLz05VqagxwRVwi4yV4,54
14
+ kscale-0.0.3.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
15
+ kscale-0.0.3.dist-info/RECORD,,
@@ -0,0 +1,2 @@
1
+ [console_scripts]
2
+ kscale = kscale.store.cli:sync_main
@@ -1,15 +0,0 @@
1
- kscale/__init__.py,sha256=QvlVh4JTl3JL7jQAja76yKtT-IvF4631ASjWY1wS6AQ,22
2
- kscale/conf.py,sha256=dnO8qii7JeMPMdjfuxnFXURRM4krwzL404RnwDkyL2M,1441
3
- kscale/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
4
- kscale/store/__init__.py,sha256=sXLh7g3KC4QCFxcZGBTpG2scR7hmmBsMjq6LqRptkRg,22
5
- kscale/store/cli.py,sha256=sf6F8ZLMenNpxi9dPB6kpaUOmRdfNBcKSnAcoN7EMo0,733
6
- kscale/store/pybullet.py,sha256=2Pog9wPSyyhmhTJY6x5KhuUQgenpTPIvp_9wxOy9ZaU,7622
7
- kscale/store/urdf.py,sha256=-dp0DEMoDNIqieG9egPv_XaV7Cg0GthMI_QIYvG2Xjk,6185
8
- kscale/store/gen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
9
- kscale/store/gen/api.py,sha256=7r6KvlZyB-kWKSCixc-YOtaWPy2mn0v11d2tuLlRMSc,7700
10
- kscale-0.0.2.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
11
- kscale-0.0.2.dist-info/METADATA,sha256=fNvMO1rM9-UfnHjp3BZ4t4FmcQ-CRPAmBDAdXRkBu-8,2028
12
- kscale-0.0.2.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
13
- kscale-0.0.2.dist-info/entry_points.txt,sha256=Vfj9O643497OONpgwy5UJLJ5a5Q1CfHZSuYKDB9D_GI,49
14
- kscale-0.0.2.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
15
- kscale-0.0.2.dist-info/RECORD,,
@@ -1,2 +0,0 @@
1
- [console_scripts]
2
- kscale = kscale.store.cli:main
File without changes