Update app_logic.py

app_logic.py

@@ -7,7 +7,6 @@ import logging
 from pathlib import Path
 from PIL import Image
 
-# Import the extractor logic
 from keylock import core as keylock_core
 
 logging.basicConfig(
@@ -16,23 +15,41 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# --- NEW FUNCTION TO GET USERNAME ---
+def get_username_from_token(api_token: str):
+    """
+    Fetches the username associated with a Hugging Face API token.
+    Returns the username on success, or an empty string on failure.
+    """
+    if not api_token:
+        logger.info("API token is empty, cannot fetch username.")
+        return ""
+
+    try:
+        logger.info("Attempting to fetch username from provided API token.")
+        user_info = whoami(token=api_token)
+        username = user_info.get('name')
+        
+        if username:
+            logger.info(f"Successfully fetched username: {username}")
+            return username
+        else:
+            logger.warning("Token was valid, but no username found in whoami() response.")
+            return "" # Return empty if name key doesn't exist
+    except Exception as e:
+        logger.error(f"Failed to validate token and fetch username: {e}")
+        return ""
+
 def _get_api_token(ui_token_from_textbox=None):
-    """Helper to get the API token from UI input or environment variables."""
     env_token = os.getenv('HF_TOKEN')
     if env_token: return env_token, None
     if ui_token_from_textbox: return ui_token_from_textbox, None
     return None, "Error: Hugging Face API token not provided."
 
 def extract_markdown_from_image(repo_image: Image.Image, image_password: str):
-    """
-    A wrapper function to call the keylock core extractor.
-    Returns the markdown string on success, or an error message on failure.
-    """
     if not repo_image:
         return None, "Error: Please upload a repository image."
-    
     password_to_use = image_password if image_password is not None else ""
-        
     try:
         logger.info("Starting extraction of markdown from image...")
         result = keylock_core.extract_repo_from_image(repo_image, password_to_use)
@@ -46,16 +63,10 @@ def extract_markdown_from_image(repo_image: Image.Image, image_password: str):
         return None, f"Error during extraction: {e}"
 
 def parse_markdown(markdown_input):
-    """
-    Parses the repo markdown string to get file structure and content.
-    This function remains the same.
-    """
     space_info = {"files": []}
     current_file_path = None; current_file_content_lines = []
     in_file_definition = False; in_code_block = False
-
     lines = markdown_input.strip().split("\n")
-    
     for line_content_orig in lines:
         line_content_stripped = line_content_orig.strip()
         if line_content_stripped.startswith("### File:"):
@@ -64,23 +75,19 @@ def parse_markdown(markdown_input):
             current_file_path = line_content_stripped.replace("### File:", "").strip()
             current_file_content_lines = []; in_file_definition = True; in_code_block = False
             continue
-        if not in_file_definition:
-            continue
+        if not in_file_definition: continue
         if line_content_stripped.startswith("```"):
             in_code_block = not in_code_block
             continue
         current_file_content_lines.append(line_content_orig)
     if current_file_path and in_file_definition:
         space_info["files"].append({"path": current_file_path, "content": "\n".join(current_file_content_lines)})
-    
     space_info["files"] = [f for f in space_info["files"] if f.get("path")]
     return space_info
 
 def _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui):
-    """Determines the final repo_id (owner/space_name)."""
     if not space_name_ui: return None, "Error: Space Name cannot be empty."
     if "/" in space_name_ui: return None, "Error: Space Name should not contain '/'. Use the Owner field."
-    
     final_owner = owner_ui
     if not final_owner:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
@@ -90,25 +97,17 @@ def _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui):
             final_owner = user_info.get('name')
         except Exception as e:
             return None, f"Error retrieving username from token: {e}. Please specify the Owner manually."
-            
     if not final_owner: return None, "Error: Owner could not be determined."
     return f"{final_owner}/{space_name_ui}", None
 
 def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, markdown_input):
-    """
-    The core function to create the space from the markdown string.
-    This function remains largely the same.
-    """
     repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
     if err: return err
-
     try:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
         if token_err: return token_err
-        
         space_info = parse_markdown(markdown_input)
         if not space_info["files"]: return "Error: No files were found after parsing the markdown from the image."
-
         with tempfile.TemporaryDirectory() as temp_dir:
             repo_staging_path = Path(temp_dir) / "repo_staging"
             repo_staging_path.mkdir()
@@ -117,11 +116,7 @@ def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, mar
                 file_path_abs.parent.mkdir(parents=True, exist_ok=True)
                 with open(file_path_abs, "w", encoding="utf-8") as f:
                     f.write(file_info["content"])
-            
-            # Create the repo
             create_repo(repo_id=repo_id, token=resolved_api_token, repo_type="space", space_sdk=sdk_ui, exist_ok=True)
-            
-            # Upload the contents
             upload_folder(
                 repo_id=repo_id,
                 folder_path=str(repo_staging_path),
@@ -130,21 +125,14 @@ def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, mar
                 repo_type="space",
                 commit_message="Initial Space creation from repo-embedded image"
             )
-            
             return f"**Success!** Space created or updated: [{repo_id}](https://huggingface.co/spaces/{repo_id})"
     except Exception as e:
         logger.exception(f"Error in create_space for {repo_id}:")
         return f"Error during Space creation: {str(e)}"
 
 def build_space_from_image(api_token, repo_image, image_password, new_space_name, owner, sdk):
-    """
-    The main orchestrator function called by the Gradio UI.
-    """
-    # 1. Extract markdown from the image
     markdown_string, error = extract_markdown_from_image(repo_image, image_password)
     if error:
-        return error # Return extraction error to the UI
-
-    # 2. Create the space using the extracted markdown
+        return error
     result = create_space(api_token, new_space_name, owner, sdk, markdown_string)
     return result