-
Notifications
You must be signed in to change notification settings - Fork 75
/
family.py
114 lines (83 loc) · 2.9 KB
/
family.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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
"""
family
======
This module defines the Family class which represents
families of particles (e.g. dm, gas, star particles).
New Family objects are automatically registered so that
snapshots can use them in the normal syntax (snap.dm,
snap.star, etc).
In practice the easiest way to make use of the flexibility
this module provides is through adding more families of
particles in your config.ini.
"""
import functools
import sys
from . import config_parser
_registry = []
def family_names(with_aliases=False):
"""Returns a list of the names of all particle families.
If with_aliases is True, include aliases in the list."""
global _registry
l = []
for o in _registry:
l.append(o.name)
if with_aliases:
for a in o.aliases:
l.append(a)
return l
def get_family(name, create=False):
"""Returns a family corresponding to the specified string. If the
family does not exist and create is False, raises ValueError. If
the family does not exist and create is True, an appropriate
object is instantiated, registered and returned."""
if isinstance(name, Family):
return name
name = name.lower()
# or should it check and raise rather than just convert?
# Not sure.
for n in _registry:
if n.name == name or name in n.aliases:
return n
if create:
return Family(name)
else:
raise ValueError(name +
" not a family") # is ValueError the right thing here?
class Family:
def __init__(self, name, aliases=[]):
if name != name.lower():
raise ValueError("Family names must be lower case")
if name in family_names(with_aliases=True):
raise ValueError("Family name " + name + " is not unique")
for a in aliases:
if a != a.lower():
raise ValueError("Aliases must be lower case")
self.name = name
self.aliases = aliases
_registry.append(self)
def __repr__(self):
return "<Family " + self.name + ">"
def __reduce__(self):
return get_family, (self.name, True), {"aliases": self.aliases}
def __iter__(self):
# Provided so a single family can be treated as a list of families
yield self
def __str__(self):
return self.name
def __cmp__(self, other):
# for python 2.x
return cmp(str(self), str(other))
def __eq__(self, other):
return str(self) == str(other)
def __lt__(self, other):
return str(self) < str(other)
def __hash__(self):
return hash(str(self))
if sys.version_info[0] >= 3:
Family = functools.total_ordering(Family)
# Instantiate the default families as specified
# by the configuration file
g = globals()
for f in config_parser.options('families'):
aliases = config_parser.get('families', f)
g[f] = Family(f, list(map(str.strip, aliases.split(","))))