-
Notifications
You must be signed in to change notification settings - Fork 62
/
mpid.py
76 lines (55 loc) · 2.12 KB
/
mpid.py
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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import re
from dataclasses import dataclass
from typing import Tuple, Union
NOTHING = object()
mpid_regex = re.compile(r"^([A-Za-z]*-)?(\d+)(-[A-Za-z0-9]+)*$")
class MPID:
"""
A Materials Project type ID with a prefix and an integer
This class enables seemlessly mixing MPIDs and regular integer IDs
Prefixed IDs are considered less than non-prefixed IDs to enable proper
mixing with the Materials Project
"""
def __init__(self, val: Union["MPID", int, str]):
if isinstance(val, MPID):
self.parts = val.parts # type: ignore
elif isinstance(val, int):
self.parts = (NOTHING, val)
elif isinstance(val, str):
parts = val.split("-")
parts[1] = int(parts[1]) # type: ignore
self.parts = tuple(parts)
else:
raise ValueError("Must provide an MPID, int, or string of the format prefix-number")
def __eq__(self, other: object):
if isinstance(other, MPID):
return self.parts == other.parts
elif isinstance(other, (int, str)):
return self.parts == MPID(other).parts
def __str__(self):
return "-".join((str(i) for i in self.parts))
def __lt__(self, other: Union["MPID", int, str]):
# Always sort MPIDs before pure integer IDs
if isinstance(other, int):
return True
elif isinstance(other, str):
other_parts = other.split("-")
other_parts[1] = int(other_parts[1]) # type: ignore
else:
other_parts = other.parts
return self.parts < other_parts
def __hash__(self):
return hash(self.parts)
@classmethod
def __get_validators__(cls):
yield cls.validate
@classmethod
def __modify_schema__(cls, field_schema):
field_schema.update(pattern=r"^([A-Za-z]*-)?(\d+)(-[A-Za-z0-9]+)*$", examples=["mp-3534", "3453", "mp-834-Ag"])
@classmethod
def validate(cls, v):
if isinstance(v, MPID):
return v
elif isinstance(v, str) and mpid_regex.fullmatch(v):
return v
raise ValueError("Invalid MPID Format")