From fa1bc9dbcbd93c3de38775dc21da9ba1271e83f8 Mon Sep 17 00:00:00 2001
From: Lucain Pouget <lucainp@gmail.com>
Date: Thu, 5 Dec 2024 17:11:30 +0100
Subject: [PATCH] Fix autocompletion not working with ModelHubMixin

---
 src/huggingface_hub/hub_mixin.py |  2 +-
 tests/test_hub_mixin.py          | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/huggingface_hub/hub_mixin.py b/src/huggingface_hub/hub_mixin.py
index cc352ff799..cbcdd74b90 100644
--- a/src/huggingface_hub/hub_mixin.py
+++ b/src/huggingface_hub/hub_mixin.py
@@ -274,7 +274,7 @@ def __init_subclass__(
         }
         cls._hub_mixin_inject_config = "config" in inspect.signature(cls._from_pretrained).parameters
 
-    def __new__(cls, *args, **kwargs) -> "ModelHubMixin":
+    def __new__(cls: Type[T], *args, **kwargs) -> T:
         """Create a new instance of the class and handle config.
 
         3 cases:
diff --git a/tests/test_hub_mixin.py b/tests/test_hub_mixin.py
index 4247443e36..10f12f7a25 100644
--- a/tests/test_hub_mixin.py
+++ b/tests/test_hub_mixin.py
@@ -7,6 +7,7 @@
 from typing import Dict, Optional, Union
 from unittest.mock import Mock, patch
 
+import jedi
 import pytest
 
 from huggingface_hub import HfApi, hf_hub_download
@@ -452,3 +453,24 @@ def test_inherited_class(self):
         model = DummyModelInherited()
         assert model._hub_mixin_info.repo_url == "https://hf.co/my-repo"
         assert model._hub_mixin_info.model_card_data.library_name == "my-cool-library"
+
+    def test_autocomplete_works_as_expected(self):
+        """Regression test for #2694.
+
+        Ensure that autocomplete works as expected when inheriting from `ModelHubMixin`.
+
+        See https://github.com/huggingface/huggingface_hub/issues/2694.
+        """
+        source = """
+from huggingface_hub import ModelHubMixin
+
+class Dummy(ModelHubMixin):
+    def dummy_example_for_test(self, x: str) -> str:
+        return x
+
+a = Dummy()
+a.dum""".strip()
+        script = jedi.Script(source, path="example.py")
+        source_lines = source.split("\n")
+        completions = script.complete(len(source_lines), len(source_lines[-1]))
+        assert any(completion.name == "dummy_example_for_test" for completion in completions)