xiaoshiai-hub 0.1.0__py3-none-any.whl

tests/__init__.py

@@ -0,0 +1,4 @@
+"""
+Tests for XiaoShi AI Hub Python SDK
+"""
+

tests/test_download.py

@@ -0,0 +1,142 @@
+"""
+Tests for download functionality
+"""
+
+import os
+import tempfile
+import unittest
+from unittest.mock import Mock, patch, MagicMock
+
+from xiaoshiai_hub.download import (
+    _match_pattern,
+    _should_download_file,
+    hf_hub_download,
+    snapshot_download,
+)
+
+
+class TestPatternMatching(unittest.TestCase):
+    """Test pattern matching functions."""
+    
+    def test_match_pattern_exact(self):
+        """Test exact pattern matching."""
+        self.assertTrue(_match_pattern("config.yaml", "config.yaml"))
+        self.assertFalse(_match_pattern("config.yml", "config.yaml"))
+    
+    def test_match_pattern_wildcard(self):
+        """Test wildcard pattern matching."""
+        self.assertTrue(_match_pattern("config.yaml", "*.yaml"))
+        self.assertTrue(_match_pattern("model.yml", "*.yml"))
+        self.assertFalse(_match_pattern("config.txt", "*.yaml"))
+    
+    def test_match_pattern_prefix(self):
+        """Test prefix pattern matching."""
+        self.assertTrue(_match_pattern("config.yaml", "config*"))
+        self.assertTrue(_match_pattern("config_v2.yaml", "config*"))
+        self.assertFalse(_match_pattern("model.yaml", "config*"))
+    
+    def test_should_download_file_no_patterns(self):
+        """Test file download decision with no patterns."""
+        self.assertTrue(_should_download_file("config.yaml"))
+        self.assertTrue(_should_download_file("model.bin"))
+    
+    def test_should_download_file_allow_patterns(self):
+        """Test file download decision with allow patterns."""
+        allow = ["*.yaml", "*.yml"]
+        self.assertTrue(_should_download_file("config.yaml", allow_patterns=allow))
+        self.assertTrue(_should_download_file("model.yml", allow_patterns=allow))
+        self.assertFalse(_should_download_file("model.bin", allow_patterns=allow))
+    
+    def test_should_download_file_ignore_patterns(self):
+        """Test file download decision with ignore patterns."""
+        ignore = [".git*", "*.tmp"]
+        self.assertFalse(_should_download_file(".gitignore", ignore_patterns=ignore))
+        self.assertFalse(_should_download_file("temp.tmp", ignore_patterns=ignore))
+        self.assertTrue(_should_download_file("config.yaml", ignore_patterns=ignore))
+    
+    def test_should_download_file_both_patterns(self):
+        """Test file download decision with both allow and ignore patterns."""
+        allow = ["*.yaml", "*.yml"]
+        ignore = [".git*"]
+        
+        self.assertTrue(_should_download_file("config.yaml", allow, ignore))
+        self.assertFalse(_should_download_file(".gitignore", allow, ignore))
+        self.assertFalse(_should_download_file("model.bin", allow, ignore))
+
+
+class TestDownloadFunctions(unittest.TestCase):
+    """Test download functions."""
+    
+    @patch('xiaoshiai_hub.download.HubClient')
+    def test_hf_hub_download(self, mock_client_class):
+        """Test hf_hub_download function."""
+        # Setup mock
+        mock_client = Mock()
+        mock_client_class.return_value = mock_client
+        
+        mock_repo_info = Mock()
+        mock_repo_info.default_branch = "main"
+        mock_client.get_repository_info.return_value = mock_repo_info
+        
+        # Test download
+        with tempfile.TemporaryDirectory() as tmpdir:
+            result = hf_hub_download(
+                repo_id="demo/demo",
+                filename="config.yaml",
+                local_dir=tmpdir,
+                username="test",
+                password="test",
+            )
+            
+            # Verify client was created with correct params
+            mock_client_class.assert_called_once()
+            
+            # Verify download was called
+            mock_client.download_file.assert_called_once()
+            
+            # Verify result path
+            self.assertIn("config.yaml", result)
+    
+    @patch('xiaoshiai_hub.download.HubClient')
+    def test_snapshot_download(self, mock_client_class):
+        """Test snapshot_download function."""
+        # Setup mock
+        mock_client = Mock()
+        mock_client_class.return_value = mock_client
+        
+        mock_repo_info = Mock()
+        mock_repo_info.default_branch = "main"
+        mock_client.get_repository_info.return_value = mock_repo_info
+        
+        # Mock content structure
+        mock_file = Mock()
+        mock_file.type = "file"
+        mock_file.path = "config.yaml"
+        
+        mock_content = Mock()
+        mock_content.entries = [mock_file]
+        mock_client.get_repository_content.return_value = mock_content
+        
+        # Test download
+        with tempfile.TemporaryDirectory() as tmpdir:
+            result = snapshot_download(
+                repo_id="demo/demo",
+                local_dir=tmpdir,
+                username="test",
+                password="test",
+                verbose=False,
+            )
+            
+            # Verify client was created
+            mock_client_class.assert_called_once()
+            
+            # Verify download was called
+            mock_client.download_file.assert_called()
+            
+            # Verify result path
+            self.assertEqual(result, tmpdir)
+
+
+if __name__ == '__main__':
+    unittest.main()
+

xiaoshiai_hub/__init__.py

@@ -0,0 +1,41 @@
+"""
+XiaoShi AI Hub Python SDK
+
+A Python library for interacting with XiaoShi AI Hub repositories.
+"""
+
+from .client import HubClient, DEFAULT_BASE_URL
+from .download import (
+    hf_hub_download,
+    snapshot_download,
+)
+from .exceptions import (
+    HubException,
+    RepositoryNotFoundError,
+    FileNotFoundError,
+    AuthenticationError,
+)
+from .types import (
+    Repository,
+    Ref,
+    GitContent,
+    Commit,
+)
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "HubClient",
+    "DEFAULT_BASE_URL",
+    "hf_hub_download",
+    "snapshot_download",
+    "HubException",
+    "RepositoryNotFoundError",
+    "FileNotFoundError",
+    "AuthenticationError",
+    "Repository",
+    "Ref",
+    "GitContent",
+    "Commit",
+]
+

xiaoshiai_hub/client.py

@@ -0,0 +1,259 @@
+"""
+XiaoShi AI Hub Client
+"""
+
+import base64
+import json
+import os
+from typing import List, Optional
+from urllib.parse import urljoin
+
+import requests
+
+try:
+    from tqdm.auto import tqdm
+except ImportError:
+    tqdm = None
+
+from .exceptions import (
+    AuthenticationError,
+    HTTPError,
+    RepositoryNotFoundError,
+)
+from .types import Repository, Ref, GitContent
+
+
+# 默认基础 URL，可通过环境变量 MOHA_ENDPOINT 覆盖
+DEFAULT_BASE_URL = os.environ.get(
+    "MOHA_ENDPOINT",
+    "https://rune.develop.xiaoshiai.cn/api/moha"
+)
+
+
+class HubClient:
+    """Client for interacting with XiaoShi AI Hub API."""
+
+    def __init__(
+        self,
+        base_url: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        token: Optional[str] = None,
+    ):
+        """
+        Initialize the Hub client.
+
+        Args:
+            base_url: Base URL of the Hub API (default: from MOHA_ENDPOINT env var)
+            username: Username for authentication
+            password: Password for authentication
+            token: Token for authentication (alternative to username/password)
+        """
+        self.base_url = (base_url or DEFAULT_BASE_URL).rstrip('/')
+        self.username = username
+        self.password = password
+        self.token = token
+        self.session = requests.Session()
+        
+        # Set up authentication
+        if token:
+            self.session.headers['Authorization'] = f'Bearer {token}'
+        elif username and password:
+            auth_string = f"{username}:{password}"
+            encoded = base64.b64encode(auth_string.encode()).decode()
+            self.session.headers['Authorization'] = f'Basic {encoded}'
+    
+    def _make_request(
+        self,
+        method: str,
+        url: str,
+        **kwargs
+    ) -> requests.Response:
+        """Make an HTTP request with error handling."""
+        try:
+            response = self.session.request(method, url, **kwargs)
+            
+            if response.status_code == 401:
+                raise AuthenticationError("Authentication failed")
+            elif response.status_code == 404:
+                raise RepositoryNotFoundError("Resource not found")
+            elif response.status_code >= 400:
+                raise HTTPError(
+                    f"HTTP {response.status_code}: {response.reason}",
+                    status_code=response.status_code
+                )
+            
+            return response
+        except requests.RequestException as e:
+            raise HTTPError(f"Request failed: {str(e)}")
+    
+    def get_repository_info(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+    ) -> Repository:
+        """
+        Get repository information.
+        
+        Args:
+            organization: Organization name
+            repo_type: Repository type ("models" or "datasets")
+            repo_name: Repository name
+            
+        Returns:
+            Repository information
+        """
+        url = f"{self.base_url}/organizations/{organization}/{repo_type}/{repo_name}"
+        response = self._make_request("GET", url)
+        data = response.json()
+        
+        return Repository(
+            name=data.get('name', repo_name),
+            organization=organization,
+            type=repo_type,
+            default_branch=data.get('defaultBranch'),
+            description=data.get('description'),
+        )
+    
+    def get_repository_refs(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+    ) -> List[Ref]:
+        """
+        Get repository references (branches and tags).
+        
+        Args:
+            organization: Organization name
+            repo_type: Repository type ("models" or "datasets")
+            repo_name: Repository name
+            
+        Returns:
+            List of references
+        """
+        url = f"{self.base_url}/organizations/{organization}/{repo_type}/{repo_name}/refs"
+        response = self._make_request("GET", url)
+        data = response.json()
+        
+        refs = []
+        for ref_data in data:
+            refs.append(Ref(
+                name=ref_data.get('name', ''),
+                ref=ref_data.get('ref', ''),
+                fully_name=ref_data.get('fullyName', ''),
+                type=ref_data.get('type', ''),
+                hash=ref_data.get('hash', ''),
+                is_default=ref_data.get('isDefault', False),
+            ))
+        
+        return refs
+    
+    def get_repository_content(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+        branch: str,
+        path: str = "",
+    ) -> GitContent:
+        """
+        Get repository content at a specific path.
+        
+        Args:
+            organization: Organization name
+            repo_type: Repository type ("models" or "datasets")
+            repo_name: Repository name
+            branch: Branch name
+            path: Path within the repository (empty for root)
+            
+        Returns:
+            Git content information
+        """
+        if path:
+            url = f"{self.base_url}/organizations/{organization}/{repo_type}/{repo_name}/contents/{branch}/{path}"
+        else:
+            url = f"{self.base_url}/organizations/{organization}/{repo_type}/{repo_name}/contents/{branch}"
+        
+        response = self._make_request("GET", url)
+        data = response.json()
+        
+        return self._parse_git_content(data)
+    
+    def _parse_git_content(self, data: dict) -> GitContent:
+        """Parse GitContent from API response."""
+        entries = None
+        if 'entries' in data and data['entries']:
+            entries = [self._parse_git_content(entry) for entry in data['entries']]
+        
+        return GitContent(
+            name=data.get('name', ''),
+            path=data.get('path', ''),
+            type=data.get('type', 'file'),
+            size=data.get('size', 0),
+            hash=data.get('hash'),
+            content_type=data.get('contentType'),
+            content=data.get('content'),
+            content_omitted=data.get('contentOmitted', False),
+            entries=entries,
+        )
+    
+    def download_file(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+        branch: str,
+        file_path: str,
+        local_path: str,
+        show_progress: bool = True,
+    ) -> None:
+        """
+        Download a single file from the repository.
+
+        Args:
+            organization: Organization name
+            repo_type: Repository type ("models" or "datasets")
+            repo_name: Repository name
+            branch: Branch name
+            file_path: Path to the file in the repository
+            local_path: Local path to save the file
+            show_progress: Whether to show download progress bar
+        """
+        url = f"{self.base_url}/organizations/{organization}/{repo_type}/{repo_name}/resolve/{branch}/{file_path}"
+        response = self._make_request("GET", url, stream=True)
+
+        # Get file size from headers
+        total_size = int(response.headers.get('content-length', 0))
+
+        # Create parent directories if needed
+        import os
+        os.makedirs(os.path.dirname(local_path) if os.path.dirname(local_path) else '.', exist_ok=True)
+
+        # Prepare progress bar
+        progress_bar = None
+        if show_progress and tqdm is not None and total_size > 0:
+            # Get filename for display
+            filename = os.path.basename(file_path)
+            progress_bar = tqdm(
+                total=total_size,
+                unit='B',
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=filename,
+                leave=True,
+            )
+
+        # Write file with progress
+        try:
+            with open(local_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+                        if progress_bar is not None:
+                            progress_bar.update(len(chunk))
+        finally:
+            if progress_bar is not None:
+                progress_bar.close()
+