Coverage for skema/skema_py/acsets.py: 69%

1from dataclasses import dataclass 

2import json 

3from typing import List, Dict 

4 

5@dataclass(eq=True, frozen=True) 

6class Ob: 

7 name: str 

8 

9class Property: 

10 pass 

11 

12@dataclass(eq=True, frozen=True) 

13class Hom(Property): 

14 name: str 

15 dom: Ob 

16 hom: Ob 

17 

18@dataclass(eq=True, frozen=True) 

19class AttrType: 

20 name: str 

21 ty: type 

22 

23@dataclass(eq=True, frozen=True) 

24class Attr(Property): 

25 name: str 

26 dom: Ob 

27 codom: AttrType 

28 

29@dataclass(eq=True, frozen=True) 

30class Schema: 

31 """Schema for an acset""" 

32 obs: List[Ob] 

33 homs: List[Hom] 

34 attrtypes: List[AttrType] 

35 attrs: List[Attr] 

36 

37 def props_outof(self, ob: Ob) -> List[Property]: 

38 return filter(lambda f: f.dom == ob, self.homs + self.attrs) 

39 

40 def from_string(self, s: str): 

41 x = next((x for x in self.obs if x.name == s), None) 

42 if x != None: 

43 return x 

44 x = next((x for x in self.homs if x.name == s), None) 

45 if x != None: 

46 return x 

47 x = next((x for x in self.attrtypes if x.name == s), None) 

48 if x != None: 

49 return x 

50 x = next((x for x in self.attrs if x.name == s), None) 

51 if x != None: 

52 return x 

53 

54class ACSet: 

55 schema: Schema 

56 _parts: Dict[Ob, int] 

57 _subparts: Dict[Property, Dict[int, any]] 

58 

59 def __init__(self, schema: Schema): 

60 self.schema = schema 

61 self._parts = { ob:0 for ob in schema.obs } 

62 self._subparts = { f:{} for f in schema.homs + schema.attrs } 

63 

64 def add_parts(self, ob: Ob, n: int) -> range: 

65 assert(ob in self.schema.obs) 

66 i = self._parts[ob] 

67 self._parts[ob] += n 

68 return range(i, i+n) 

69 

70 def add_part(self, ob: Ob) -> int: 

71 return self.add_parts(ob, 1)[0] 

72 

73 def _check_type(self, f: Property, x: any): 

74 if f in self.schema.homs: 

75 assert(isinstance(x, int)) 

76 elif f in self.schema.attrs: 

77 assert(isinstance(x, f.codom.ty)) 

78 else: 

79 raise(f"{f} not found in schema") 

80 

81 def set_subpart(self, i: int, f: Property, x: any): 

82 self._check_type(f, x) 

83 self._subparts[f][i] = x 

84 

85 def subpart(self, i: int, f: Property, oneindex=False): 

86 if oneindex and type(f) == Hom: 

87 return self._subparts[f][i] + 1 

88 else: 

89 return self._subparts[f].get(i, None) 

90 

91 def nparts(self, ob: Ob) -> int: 

92 assert(ob in self.schema.obs) 

93 return self._parts[ob] 

94 

95 def parts(self, ob: Ob) -> range: 

96 return range(0, self.nparts(ob)) 

97 

98 def incident(self, x: any, f: Property) -> List[int]: 

99 self._check_type(f, x) 

100 return filter(lambda i: self.subpart(i, f) == x, self.parts(f.dom)) 

101 

102 def write_json(self): 

103 print(self._subparts) 

104 return json.dumps({ 

105 ob.name: [self.prop_dict(ob, i) for i in self.parts(ob)] for ob in self.schema.obs 

106 }) 

107 

108 def prop_dict(self, ob: Ob, i: int) -> Dict[str, any]: 

109 return { f.name: self.subpart(i, f, oneindex=True) for f in self.schema.props_outof(ob) } 

110 

111 @classmethod 

112 def read_json(cls, schema: Schema, s: str): 

113 acs = cls(schema) 

114 d = json.loads(s) 

115 for (obname, proplist) in d.items(): 

116 ob = schema.from_string(obname) 

117 for props in proplist: 

118 i = acs.add_part(ob) 

119 for (fname,v) in props.items(): 

120 f = schema.from_string(fname) 

121 if type(f) == Hom: 

122 acs.set_subpart(i, f, v-1) 

123 else: 

124 acs.set_subpart(i, f, v) 

125 return acs