storage verison added

modules/storage.py

@@ -1,4 +1,5 @@
 # modules/storage.py
+__version__ = "0.1.0" # Added version
 import os
 import urllib.parse
 import tempfile
@@ -106,10 +107,17 @@ def upload_files_to_repo(
             valid_files.append(f)
     
     if not valid_files:
-        return []  # or raise an exception
+        # Return a dictionary with None values for permalinks if create_permalink was True
+        if create_permalink:
+            return {
+                "response": "No valid files to upload.",
+                "permalink": None,
+                "short_permalink": None
+            }
+        return [] 
     
     # Create a temporary directory; copy valid files directly into it.
-    with tempfile.TemporaryDirectory() as temp_dir:
+    with tempfile.TemporaryDirectory(dir=os.getenv("TMPDIR", "/tmp")) as temp_dir:
         for file_path in valid_files:
             filename = os.path.basename(file_path)
             dest_path = os.path.join(temp_dir, filename)
@@ -135,22 +143,32 @@ def upload_files_to_repo(
     
     # If permalink creation is requested and exactly 3 valid files are provided,
     # try to generate a permalink using generate_permalink().
-    if create_permalink and len(valid_files) == 3:
+    if create_permalink: # No need to check len(valid_files) == 3 here, generate_permalink will handle it
         permalink = generate_permalink(valid_files, base_url_external, permalink_viewer_url)
         if permalink:
-            result, short_id = gen_full_url(
+            status, short_id = gen_full_url(
                 full_url=permalink,
-                repo_id=HF_REPO_ID,
-                json_file=SHORTENER_JSON_FILE
+                repo_id=HF_REPO_ID, # This comes from constants
+                json_file=SHORTENER_JSON_FILE # This comes from constants
             )
-            permalink_short = f"https://{permalink_viewer_url}/?sid={short_id}"
-            print(f"Creating shortened URL: {result} - {short_id}")
+            if status in ["created_short", "success_retrieved_short", "exists_match"]:
+                permalink_short = f"https://{permalink_viewer_url}/?sid={short_id}"
+            else: # Shortening failed or conflict not resolved to a usable short_id
+                permalink_short = None 
+                print(f"URL shortening status: {status} for {permalink}")
+
             return {
                 "response": response,
                 "permalink": permalink,
                 "short_permalink": permalink_short
             }
-    
+        else: # generate_permalink returned None (criteria not met)
+            return {
+                "response": response, # Still return upload response
+                "permalink": None,
+                "short_permalink": None
+            }
+
     # Otherwise, return individual tuples for each file.
     return [(response, link) for link in individual_links]
 
@@ -190,7 +208,11 @@ def _upload_json_to_repo(data, repo_id, json_file_name, repo_type="dataset"):
     try:
         login(token=HF_API_TOKEN)
         api = HfApi()
-        with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json") as tmp_file:
+        # Use a temporary directory specified by TMPDIR or default to system temp
+        temp_dir_for_json = os.getenv("TMPDIR", tempfile.gettempdir())
+        os.makedirs(temp_dir_for_json, exist_ok=True)
+
+        with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json", dir=temp_dir_for_json) as tmp_file:        
             json.dump(data, tmp_file, indent=2)
             tmp_file_path = tmp_file.name
         
@@ -229,7 +251,7 @@ def _find_url_in_json(data, short_url=None, full_url=None):
 
 def _add_url_to_json(data, short_url, full_url):
     """Adds a new short_url/full_url pair to the data. Returns updated data."""
-    if data is None: # Initialize if data is None
+    if data is None:
         data = []
     data.append({"short_url": short_url, "full_url": full_url})
     return data
@@ -259,10 +281,7 @@ def gen_full_url(short_url=None, full_url=None, repo_id=None, repo_type="dataset
     # Case 1: Only short_url provided (lookup full_url)
     if short_url and not full_url:
         found_full_url = _find_url_in_json(url_data, short_url=short_url)
-        if found_full_url:
-            return "success_retrieved_full", found_full_url
-        else:
-            return "not_found_short", None
+        return ("success_retrieved_full", found_full_url) if found_full_url else ("not_found_short", None)
 
     # Case 2: Only full_url provided (lookup or create short_url)
     if full_url and not short_url:
@@ -272,9 +291,6 @@ def gen_full_url(short_url=None, full_url=None, repo_id=None, repo_type="dataset
         else:
             # Create new short_url
             new_short_id = _generate_short_id()
-            # Construct the short URL using the permalink_viewer_url and the new_short_id as a query parameter or path
-            # For this example, let's assume short URLs are like: https://permalink_viewer_url/?id=new_short_id
-            
             url_data = _add_url_to_json(url_data, new_short_id, full_url)
             if _upload_json_to_repo(url_data, repo_id, json_file, repo_type):
                 return "created_short", new_short_id 
@@ -286,25 +302,20 @@ def gen_full_url(short_url=None, full_url=None, repo_id=None, repo_type="dataset
         found_full_for_short = _find_url_in_json(url_data, short_url=short_url)
         found_short_for_full = _find_url_in_json(url_data, full_url=full_url)
 
-        if found_full_for_short and found_full_for_short == full_url: # Both exist and match
+        if found_full_for_short == full_url: 
             return "exists_match", short_url 
-        elif found_short_for_full and found_short_for_full == short_url: # Both exist and match (redundant check, but safe)
-             return "exists_match", short_url
-        elif found_full_for_short and found_full_for_short != full_url: # short_url exists but maps to a different full_url
-            return "error_conflict_short_exists", short_url # Or perhaps update, depending on desired behavior
-        elif found_short_for_full and found_short_for_full != short_url: # full_url exists but maps to a different short_url
-             # This implies the user provided a short_url that is *not* the one already associated with this full_url.
-             # We should probably return the *existing* short_url for that full_url.
-            return "exists_full_maps_to_different_short", found_short_for_full
-        else: # Neither exists, or one exists but not the pair. Create new entry.
-            # Check if the provided short_url is already in use by another full_url
-            if _find_url_in_json(url_data, short_url=short_url) is not None:
-                return "error_conflict_short_id_taken", short_url
-
-            url_data = _add_url_to_json(url_data, short_url, full_url)
-            if _upload_json_to_repo(url_data, repo_id, json_file, repo_type):
-                return "created_specific_pair", short_url
-            else:
-                return "error_upload", None
+        if found_full_for_short is not None and found_full_for_short != full_url: 
+            return "error_conflict_short_exists_different_full", short_url
+        if found_short_for_full is not None and found_short_for_full != short_url:
+            return "error_conflict_full_exists_different_short", found_short_for_full
+        
+        # If short_url is provided and not found, or full_url is provided and not found,
+        # or neither is found, then create a new entry with the provided short_url and full_url.
+        # This effectively allows specifying a custom short_url if it's not already taken.
+        url_data = _add_url_to_json(url_data, short_url, full_url)
+        if _upload_json_to_repo(url_data, repo_id, json_file, repo_type):
+            return "created_specific_pair", short_url
+        else:
+            return "error_upload", None
                 
     return "error_unhandled_case", None # Should not be reached