Update user_history.py

user_history.py

@@ -1,6 +1,5 @@
 """
 User History is a plugin that you can add to your Spaces to cache generated images for your users.
-
 Key features:
 - 🤗 Sign in with Hugging Face
 - Save generated images with their metadata: prompts, timestamp, hyper-parameters, etc.
@@ -8,7 +7,6 @@ Key features:
 - Delete your history to respect privacy.
 - Compatible with Persistent Storage for long-term storage.
 - Admin panel to check configuration and disk usage .
-
 Useful links:
 - Demo: https://huggingface.co/spaces/Wauplin/gradio-user-history
 - README: https://huggingface.co/spaces/Wauplin/gradio-user-history/blob/main/README.md
@@ -37,18 +35,13 @@ def setup(folder_path: str | Path | None = None) -> None:
     user_history.folder_path = _resolve_folder_path(folder_path)
     user_history.initialized = True
 
-    # TODO: remove this section once all Spaces have migrated
-    _migrate_history()
-
 
 def render() -> None:
     user_history = _UserHistory()
 
     # initialize with default config
     if not user_history.initialized:
-        print(
-            "Initializing user history with default config. Use `user_history.setup(...)` to customize folder_path."
-        )
+        print("Initializing user history with default config. Use `user_history.setup(...)` to customize folder_path.")
         setup()
 
     # Render user history tab
@@ -83,18 +76,11 @@ def render() -> None:
 
     # "Export zip" row (hidden by default)
     with gr.Row():
-        export_file = gr.File(
-            file_count="single",
-            file_types=[".zip"],
-            label="Exported history",
-            visible=False,
-        )
+        export_file = gr.File(file_count="single", file_types=[".zip"], label="Exported history", visible=False)
 
     # "Config deletion" row (hidden by default)
     with gr.Row():
-        confirm_button = gr.Button(
-            "Confirm delete all history", variant="stop", visible=False
-        )
+        confirm_button = gr.Button("Confirm delete all history", variant="stop", visible=False)
         cancel_button = gr.Button("Cancel", visible=False)
 
     # Gallery
@@ -117,12 +103,8 @@ def render() -> None:
     gallery.attach_load_event(_fetch_user_history, every=None)
 
     # Interactions
-    refresh_button.click(
-        fn=_fetch_user_history, inputs=[], outputs=[gallery], queue=False
-    )
-    export_button.click(
-        fn=_export_user_history, inputs=[], outputs=[export_file], queue=False
-    )
+    refresh_button.click(fn=_fetch_user_history, inputs=[], outputs=[gallery], queue=False)
+    export_button.click(fn=_export_user_history, inputs=[], outputs=[export_file], queue=False)
 
     # Taken from https://github.com/gradio-app/gradio/issues/3324#issuecomment-1446382045
     delete_button.click(
@@ -203,9 +185,7 @@ class _UserHistory(object):
 
     def _user_lock(self, username: str) -> FileLock:
         """Ensure history is not corrupted if concurrent calls."""
-        return FileLock(
-            self.folder_path / f"{username}.lock"
-        )  # lock outside of folder => better when exporting ZIP
+        return FileLock(self.folder_path / f"{username}.lock")  # lock outside of folder => better when exporting ZIP
 
     def _user_jsonl_path(self, username: str) -> Path:
         return self._user_path(username) / "history.jsonl"
@@ -225,9 +205,7 @@ def _fetch_user_history(profile: gr.OAuthProfile | None) -> List[Tuple[str, str]
 
     user_history = _UserHistory()
     if not user_history.initialized:
-        warnings.warn(
-            "User history is not set in Gradio demo. You must use `user_history.render(...)` first."
-        )
+        warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
         return []
 
     with user_history._user_lock(username):
@@ -253,17 +231,13 @@ def _export_user_history(profile: gr.OAuthProfile | None) -> Dict | None:
 
     user_history = _UserHistory()
     if not user_history.initialized:
-        warnings.warn(
-            "User history is not set in Gradio demo. You must use `user_history.render(...)` first."
-        )
+        warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
         return None
 
     # Zip history
     with user_history._user_lock(username):
         path = shutil.make_archive(
-            str(_archives_path() / f"history_{username}"),
-            "zip",
-            user_history._user_path(username),
+            str(_archives_path() / f"history_{username}"), "zip", user_history._user_path(username)
         )
 
     return gr.update(visible=True, value=path)
@@ -278,9 +252,7 @@ def _delete_user_history(profile: gr.OAuthProfile | None) -> None:
 
     user_history = _UserHistory()
     if not user_history.initialized:
-        warnings.warn(
-            "User history is not set in Gradio demo. You must use `user_history.render(...)` first."
-        )
+        warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
         return
 
     with user_history._user_lock(username):
@@ -317,9 +289,7 @@ def _resolve_folder_path(folder_path: str | Path | None) -> Path:
     if folder_path is not None:
         return Path(folder_path).expanduser().resolve()
 
-    if os.getenv("SYSTEM") == "spaces" and os.path.exists(
-        "/data"
-    ):  # Persistent storage is enabled!
+    if os.getenv("SYSTEM") == "spaces" and os.path.exists("/data"):  # Persistent storage is enabled!
         return Path("/data") / "_user_history"
 
     # Not in a Space or Persistent storage not enabled => local folder
@@ -357,21 +327,13 @@ def _display_if_admin() -> Callable:
 def _admin_content() -> str:
     return f"""
 ## Admin section
-
 Running on **{os.getenv("SYSTEM", "local")}** (id: {os.getenv("SPACE_ID")}). {_get_msg_is_persistent_storage_enabled()}
-
 Admins: {', '.join(_fetch_admins())}
-
 {_get_nb_users()} user(s), {_get_nb_images()} image(s)
-
 ### Configuration
-
 History folder: *{_UserHistory().folder_path}*
-
 Exports folder: *{_archives_path()}*
-
 ### Disk usage
-
 {_disk_space_warning_message()}
 """
 
@@ -380,10 +342,8 @@ def _get_nb_users() -> int:
     user_history = _UserHistory()
     if not user_history.initialized:
         return 0
-    if user_history.folder_path is not None:
-        return len(
-            [path for path in user_history.folder_path.iterdir() if path.is_dir()]
-        )
+    if user_history.folder_path is not None and user_history.folder_path.exists():
+        return len([path for path in user_history.folder_path.iterdir() if path.is_dir()])
     return 0
 
 
@@ -391,7 +351,7 @@ def _get_nb_images() -> int:
     user_history = _UserHistory()
     if not user_history.initialized:
         return 0
-    if user_history.folder_path is not None:
+    if user_history.folder_path is not None and user_history.folder_path.exists():
         return len([path for path in user_history.folder_path.glob("*/images/*")])
     return 0
 
@@ -425,14 +385,10 @@ def _disk_space_warning_message() -> str:
 
 
 def _get_disk_usage(path: Path) -> Tuple[int, int, int]:
-    for path in [path] + list(
-        path.parents
-    ):  # first check target_dir, then each parents one by one
+    for path in [path] + list(path.parents):  # first check target_dir, then each parents one by one
         try:
             return shutil.disk_usage(path)
-        except (
-            OSError
-        ):  # if doesn't exist or can't read => fail silently and try parent one
+        except OSError:  # if doesn't exist or can't read => fail silently and try parent one
             pass
     return 0, 0, 0
 
@@ -451,74 +407,7 @@ def _fetch_admins() -> List[str]:
     # Running in Space => try to fetch organization members
     # Otherwise, it's not an organization => namespace is the user
     namespace = space_id.split("/")[0]
-    response = requests.get(
-        f"https://huggingface.co/api/organizations/{namespace}/members"
-    )
+    response = requests.get(f"https://huggingface.co/api/organizations/{namespace}/members")
     if response.status_code == 200:
-        return sorted(
-            (member["user"] for member in response.json()), key=lambda x: x.lower()
-        )
+        return sorted((member["user"] for member in response.json()), key=lambda x: x.lower())
     return [namespace]
-
-
-################################################################
-# Legacy helpers to migrate image structure to new data format #
-################################################################
-# TODO: remove this section once all Spaces have migrated
-
-
-def _migrate_history():
-    """Script to migrate user history from v0 to v1."""
-    legacy_history_path = _legacy_get_history_folder_path()
-    if not legacy_history_path.exists():
-        return
-
-    error_count = 0
-    for json_path in legacy_history_path.glob("*.json"):
-        username = json_path.stem
-        print(f"Migrating history for user {username}...")
-        error_count += _legacy_move_user_history(username)
-        print("Done.")
-    print(f"Migration complete. {error_count} error(s) happened.")
-
-    if error_count == 0:
-        shutil.rmtree(legacy_history_path, ignore_errors=True)
-
-
-def _legacy_move_user_history(username: str) -> int:
-    history = _legacy_read_user_history(username)
-    error_count = 0
-    for image, prompt in reversed(history):
-        try:
-            save_image(
-                label=prompt, image=image, profile={"preferred_username": username}
-            )
-        except Exception as e:
-            print("Issue while migrating image:", e)
-            error_count += 1
-    return error_count
-
-
-def _legacy_get_history_folder_path() -> Path:
-    _folder = os.environ.get("HISTORY_FOLDER")
-    if _folder is None:
-        _folder = Path(__file__).parent / "history"
-    return Path(_folder)
-
-
-def _legacy_read_user_history(username: str) -> List[Tuple[str, str]]:
-    """Return saved history for that user."""
-    with _legacy_user_lock(username):
-        path = _legacy_user_history_path(username)
-        if path.exists():
-            return json.loads(path.read_text())
-        return []  # No history yet
-
-
-def _legacy_user_history_path(username: str) -> Path:
-    return _legacy_get_history_folder_path() / f"{username}.json"
-
-
-def _legacy_user_lock(username: str) -> FileLock:
-    """Ensure history is not corrupted if concurrent calls."""
-    return FileLock(f"{_legacy_user_history_path(username)}.lock")