support new syntax in python client

python-client/wmill/wmill/client.py

@@ -11,7 +11,8 @@ import time
 import warnings
 import json
 from json import JSONDecodeError
-from typing import Dict, Any, Union, Literal
+from typing import Dict, Any, Union, Literal, Optional
+import re
 
 import httpx
 
@@ -312,6 +313,7 @@ class Windmill:
             return result_text
 
     def get_variable(self, path: str) -> str:
+        path = parse_variable_syntax(path) or path
         if self.mocked_api is not None:
             variables = self.mocked_api["variables"]
             try:
@@ -326,6 +328,7 @@ class Windmill:
         return self.get(f"/w/{self.workspace}/variables/get_value/{path}").json()
 
     def set_variable(self, path: str, value: str, is_secret: bool = False) -> None:
+        path = parse_variable_syntax(path) or path
         if self.mocked_api is not None:
             self.mocked_api["variables"][path] = value
             return
@@ -358,6 +361,7 @@ class Windmill:
         path: str,
         none_if_undefined: bool = False,
     ) -> dict | None:
+        path = parse_resource_syntax(path) or path
         if self.mocked_api is not None:
             resources = self.mocked_api["resources"]
             try:
@@ -391,6 +395,7 @@ class Windmill:
         path: str,
         resource_type: str,
     ):
+        path = parse_resource_syntax(path) or path
         if self.mocked_api is not None:
             self.mocked_api["resources"][path] = value
             return
@@ -485,6 +490,7 @@ class Windmill:
         Convenient helpers that takes an S3 resource as input and returns the settings necessary to
         initiate an S3 connection from DuckDB
         """
+        s3_resource_path = parse_resource_syntax(s3_resource_path) or s3_resource_path
         try:
             raw_obj = self.post(
                 f"/w/{self.workspace}/job_helpers/v2/duckdb_connection_settings",
@@ -506,6 +512,7 @@ class Windmill:
         Convenient helpers that takes an S3 resource as input and returns the settings necessary to
         initiate an S3 connection from Polars
         """
+        s3_resource_path = parse_resource_syntax(s3_resource_path) or s3_resource_path
         try:
             raw_obj = self.post(
                 f"/w/{self.workspace}/job_helpers/v2/polars_connection_settings",
@@ -527,6 +534,7 @@ class Windmill:
         Convenient helpers that takes an S3 resource as input and returns the settings necessary to
         initiate an S3 connection using boto3
         """
+        s3_resource_path = parse_resource_syntax(s3_resource_path) or s3_resource_path
         try:
             s3_resource = self.post(
                 f"/w/{self.workspace}/job_helpers/v2/s3_resource_info",
@@ -540,7 +548,7 @@ class Windmill:
                 "Could not generate Boto3 S3 connection settings from the provided resource"
             ) from e
 
-    def load_s3_file(self, s3object: S3Object, s3_resource_path: str | None) -> bytes:
+    def load_s3_file(self, s3object: S3Object | str, s3_resource_path: str | None) -> bytes:
         """
         Load a file from the workspace s3 bucket and returns its content as bytes.
 
@@ -552,11 +560,12 @@ class Windmill:
         file_content = my_obj_content.decode("utf-8")
         '''
         """
+        s3object = parse_s3_object(s3object)
         with self.load_s3_file_reader(s3object, s3_resource_path) as file_reader:
             return file_reader.read()
 
     def load_s3_file_reader(
-        self, s3object: S3Object, s3_resource_path: str | None
+        self, s3object: S3Object | str, s3_resource_path: str | None
     ) -> BufferedReader:
         """
         Load a file from the workspace s3 bucket and returns the bytes stream.
@@ -569,6 +578,7 @@ class Windmill:
             print(file_reader.read())
         '''
         """
+        s3object = parse_s3_object(s3object)
         reader = S3BufferedReader(
             f"{self.workspace}",
             self.client,
@@ -580,7 +590,7 @@ class Windmill:
 
     def write_s3_file(
         self,
-        s3object: S3Object | None,
+        s3object: S3Object | str | None,
         file_content: BufferedReader | bytes,
         s3_resource_path: str | None,
         content_type: str | None = None,
@@ -603,6 +613,7 @@ class Windmill:
             client.write_s3_file(s3_obj, my_file)
         '''
         """
+        s3object = parse_s3_object(s3object)
         # httpx accepts either bytes or "a bytes generator" as content. If it's a BufferedReader, we need to convert it to a generator
         if isinstance(file_content, BufferedReader):
             content_payload = bytes_generator(file_content)
@@ -644,12 +655,12 @@ class Windmill:
             raise Exception("Could not write file to S3") from e
         return S3Object(s3=response["file_key"])
 
-    def sign_s3_objects(self, s3_objects: list[S3Object]) -> list[S3Object]:
+    def sign_s3_objects(self, s3_objects: list[S3Object | str]) -> list[S3Object]:
         return self.post(
-            f"/w/{self.workspace}/apps/sign_s3_objects", json={"s3_objects": s3_objects}
+            f"/w/{self.workspace}/apps/sign_s3_objects", json={"s3_objects": list(map(parse_s3_object, s3_objects))}
         ).json()
 
-    def sign_s3_object(self, s3_object: S3Object) -> S3Object:
+    def sign_s3_object(self, s3_object: S3Object | str) -> S3Object:
         return self.post(
             f"/w/{self.workspace}/apps/sign_s3_objects",
             json={"s3_objects": [s3_object]},
@@ -1027,7 +1038,7 @@ def boto3_connection_settings(s3_resource_path: str = "") -> Boto3ConnectionSett
 
 
 @init_global_client
-def load_s3_file(s3object: S3Object, s3_resource_path: str | None = None) -> bytes:
+def load_s3_file(s3object: S3Object | str, s3_resource_path: str | None = None) -> bytes:
     """
     Load the entire content of a file stored in S3 as bytes
     """
@@ -1038,7 +1049,7 @@ def load_s3_file(s3object: S3Object, s3_resource_path: str | None = None) -> byt
 
 @init_global_client
 def load_s3_file_reader(
-    s3object: S3Object, s3_resource_path: str | None = None
+    s3object: S3Object | str, s3_resource_path: str | None = None
 ) -> BufferedReader:
     """
     Load the content of a file stored in S3
@@ -1050,7 +1061,7 @@ def load_s3_file_reader(
 
 @init_global_client
 def write_s3_file(
-    s3object: S3Object | None,
+    s3object: S3Object | str | None,
     file_content: BufferedReader | bytes,
     s3_resource_path: str | None = None,
     content_type: str | None = None,
@@ -1075,7 +1086,7 @@ def write_s3_file(
 
 
 @init_global_client
-def sign_s3_objects(s3_objects: list[S3Object]) -> list[S3Object]:
+def sign_s3_objects(s3_objects: list[S3Object | str]) -> list[S3Object]:
     """
     Sign S3 objects to be used by anonymous users in public apps
     Returns a list of signed s3 tokens
@@ -1084,7 +1095,7 @@ def sign_s3_objects(s3_objects: list[S3Object]) -> list[S3Object]:
 
 
 @init_global_client
-def sign_s3_object(s3_object: S3Object) -> S3Object:
+def sign_s3_object(s3_object: S3Object| str) -> S3Object:
     """
     Sign S3 object to be used by anonymous users in public apps
     Returns a signed s3 object
@@ -1336,3 +1347,31 @@ def task(*args, **kwargs):
         return f(args[0], None)
     else:
         return lambda x: f(x, kwargs.get("tag"))
+
+def parse_resource_syntax(s: str) -> Optional[str]:
+    """Parse resource syntax from string."""
+    if s is None:
+        return None
+    if s.startswith("$res:"):
+        return s[5:]
+    if s.startswith("res://"):
+        return s[6:]
+    return None
+
+def parse_s3_object(s3_object: S3Object | str) -> S3Object:
+    """Parse S3 object from string or S3Object format."""
+    if isinstance(a, str):
+        match = re.match(r'^s3://([^/]*)/(.*)$', s3_object)
+        if match:
+            return S3Object(s3=match.group(2) or "", storage=match.group(1) or None)
+        return S3Object(s3="")
+    else:
+        return s3_object
+
+    
+
+def parse_variable_syntax(s: str) -> Optional[str]:
+    """Parse variable syntax from string."""
+    if s.startswith("var://"):
+        return s[6:]
+    return None