forked from replicate/replicate-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathidentifier.py
More file actions
50 lines (40 loc) · 1.65 KB
/
Copy pathidentifier.py
File metadata and controls
50 lines (40 loc) · 1.65 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import re
from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Union
if TYPE_CHECKING:
from replicate.model import Model
from replicate.version import Version
class ModelVersionIdentifier(NamedTuple):
"""
A reference to a model version in the format owner/name or owner/name:version.
"""
owner: str
name: str
version: Optional[str] = None
@classmethod
def parse(cls, ref: str) -> "ModelVersionIdentifier":
"""
Split a reference in the format owner/name:version into its components.
"""
match = re.match(r"^(?P<owner>[^/]+)/(?P<name>[^/:]+)(:(?P<version>.+))?$", ref)
if not match:
raise ValueError(
f"Invalid reference to model version: {ref}. Expected format: owner/name:version"
)
return cls(match.group("owner"), match.group("name"), match.group("version"))
def _resolve(
ref: Union["Model", "Version", "ModelVersionIdentifier", str],
) -> Tuple[Optional["Version"], Optional[str], Optional[str], Optional[str]]:
from replicate.model import Model # pylint: disable=import-outside-toplevel
from replicate.version import Version # pylint: disable=import-outside-toplevel
version = None
owner, name, version_id = None, None, None
if isinstance(ref, Model):
owner, name = ref.owner, ref.name
elif isinstance(ref, Version):
version = ref
version_id = ref.id
elif isinstance(ref, ModelVersionIdentifier):
owner, name, version_id = ref
elif isinstance(ref, str):
owner, name, version_id = ModelVersionIdentifier.parse(ref)
return version, owner, name, version_id