Skip to content

Commit

Permalink
Fix autocompletion not working with ModelHubMixin
Browse files Browse the repository at this point in the history
  • Loading branch information
Wauplin committed Dec 5, 2024
1 parent 897c770 commit fa1bc9d
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 1 deletion.
2 changes: 1 addition & 1 deletion src/huggingface_hub/hub_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
22 changes: 22 additions & 0 deletions tests/test_hub_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

0 comments on commit fa1bc9d

Please sign in to comment.