@@ -48,6 +48,8 @@ def extract_orm_attr(self, text: str):
4848 not_orm = True
4949 properties = {}
5050 orm_columns = ["Column" , "Field" , "relationship" , "ForeignKey" ]
51+ pony_orm_fields = ["Required" , "Set" , "Optional" , "PrimaryKey" ]
52+ orm_columns .extend (pony_orm_fields )
5153 for i in orm_columns :
5254 if i in text :
5355 not_orm = False
@@ -57,34 +59,44 @@ def extract_orm_attr(self, text: str):
5759 text = text [index + 1 : - 1 ] # noqa E203
5860 text = text .split ("," )
5961 text = self .clean_up_cases_with_inner_pars (text )
62+ prop_index = 1
6063 if i == "Field" :
6164 _type , properties = get_django_info (text , base_text , properties )
6265 prop_index = 0
6366 elif i == "ForeignKey" :
6467 # mean it is a Django model.ForeignKey
6568 _type = "serial"
6669 properties ["foreign_key" ] = text [0 ]
67- prop_index = 1
70+ elif i in pony_orm_fields :
71+ # mean it is a Pony ORM
72+ _type , properties = get_pony_orm_info (
73+ text , i , base_text , properties
74+ )
6875 else :
69- prop_index = 1
7076 _type = text [0 ]
7177 if i == "relationship" :
7278 properties ["relationship" ] = True
73- for i in text [prop_index :]:
74- if "=" in i :
75- # can be backref=db.backref('pages', lazy=True)
76- index = i .find ("=" )
77- left = i [:index ].strip ()
78- right = i [index + 1 :].strip () # noqa: E203
79- if left == "default" :
80- default = right
81- else :
82- properties [left ] = right
83- elif "foreign" in i .lower ():
84- properties ["foreign_key" ] = i .split ("(" )[1 ].split (")" )[0 ]
79+ for item in text [prop_index :]:
80+ properties , default = self .add_property (item , properties )
8581 break
8682 return default , _type , properties , not_orm
8783
84+ @staticmethod
85+ def add_property (item : str , properties : Dict ) -> Tuple [Dict , str ]:
86+ default = None
87+ if "=" in item :
88+ # can be backref=db.backref('pages', lazy=True)
89+ index = item .find ("=" )
90+ left = item [:index ].strip ()
91+ right = item [index + 1 :].strip () # noqa: E203
92+ if left == "default" :
93+ default = right
94+ else :
95+ properties [left ] = right
96+ elif "foreign" in item .lower ():
97+ properties ["foreign_key" ] = item .split ("(" )[1 ].split (")" )[0 ]
98+ return properties , default
99+
88100 def extractor (self , text : str ) -> Dict :
89101 _type = None
90102 default = None
@@ -102,12 +114,11 @@ def visit_right_part(self, node, visited_children):
102114
103115 def visit_attr_def (self , node , visited_children ):
104116 """Makes a dict of the section (as key) and the key/value pairs."""
105- left = node .children [1 ].children [0 ].text .strip ()
117+ left = node .children [1 ].children [0 ].children [ 0 ]. text .strip ()
106118 default = None
107119 _type = None
108120 if "def " in left :
109121 attr = {"attr" : {"name" : None , "type" : _type , "default" : default }}
110-
111122 return attr
112123 if ":" in left :
113124 _type = left .split (":" )[- 1 ].strip ()
@@ -183,7 +194,6 @@ def visit_expr(self, node, visited_children):
183194 children_values [n ]["properties" ]["init" ] = final_child [
184195 "properties"
185196 ]["init" ]
186-
187197 return children_values
188198
189199 def visit_type (self , node , visited_children ):
@@ -197,17 +207,36 @@ def generic_visit(self, node, visited_children):
197207
198208
199209def process_no_name_attrs (final_child : Dict , child : Dict ) -> None :
200- if child ["attr" ]["default" ]:
201- final_child ["attrs" ][- 1 ]["default" ] = child ["attr" ]["default" ]
202- if not final_child ["attrs" ][- 1 ].get ("properties" ):
203- final_child ["attrs" ][- 1 ]["properties" ] = {}
204- elif child ["attr" ]["type" ]:
205- final_child ["attrs" ][- 1 ]["default" ] += f':{ child ["attr" ]["type" ]} '
210+ if final_child ["attrs" ]:
211+ if child ["attr" ]["default" ]:
212+ final_child ["attrs" ][- 1 ]["default" ] = child ["attr" ]["default" ]
213+ if not final_child ["attrs" ][- 1 ].get ("properties" ):
214+ final_child ["attrs" ][- 1 ]["properties" ] = {}
215+ elif child ["attr" ]["type" ] and final_child ["attrs" ][- 1 ]["default" ]:
216+ final_child ["attrs" ][- 1 ]["default" ] += f':{ child ["attr" ]["type" ]} '
206217 return final_child
207218
208219
220+ def get_pony_orm_info (
221+ text : list , field : str , base_text : str , properties : Dict
222+ ) -> Tuple :
223+ if field == "Required" :
224+ properties ["nullable" ] = False
225+ elif field == "PrimaryKey" :
226+ properties ["primary_key" ] = True
227+ elif field == "Optional" :
228+ properties ["nullable" ] = True
229+ elif field == "Set" :
230+ # relationship
231+ properties ["relationship" ] = True
232+ properties ["foreign_key" ] = text [0 ]
233+ _type = text [0 ]
234+
235+ return _type , properties
236+
237+
209238def get_django_info (text : list , base_text : str , properties : Dict ) -> Tuple :
210- # for tortoise orm & django orm
239+ # for tortoise orm & django orm
211240 split_by_field = base_text .split ("Field" )[0 ].split ("." )
212241 if len (split_by_field ) == 2 :
213242 _type = split_by_field [1 ]
0 commit comments