2828from .models import samplingfeatures
2929from .models import simulation
3030
31+ import warnings
32+ warnings .simplefilter ("ignore" , category = sqlalchemy .exc .SAWarning )
33+
34+ OUTPUT_FORMATS = ('json' , 'dataframe' , 'dict' )
35+
3136
32- class OutputFormats (Enum ):
33- JSON = 'JSON'
34- DATAFRAME = 'DATAFRAME'
35- DICT = 'DICT'
3637
3738class Base ():
3839
@@ -82,19 +83,25 @@ def __init__(self, session_maker:sqlalchemy.orm.sessionmaker) -> None:
8283
8384 def read_query (self ,
8485 query : Union [Query , Select ],
85- output_format :OutputFormats = OutputFormats . JSON ,
86+ output_format :str = 'json' ,
8687 orient :str = 'records' ) -> Union [str , pd .DataFrame ]:
88+
89+ # guard against invalid output_format strings
90+ if output_format not in OUTPUT_FORMATS :
91+ raise ValueError (f':argument output_format={ output_format } , is not a valid output_format strings: { OUTPUT_FORMATS } ' )
92+
93+ # use SQLAlchemy session to read_query and return response in the designated output_format
8794 with self .session_maker () as session :
8895 if isinstance (query , Select ):
8996 df = pd .read_sql (query , session .bind )
9097 else :
9198 df = pd .read_sql (query .statement , session .bind )
9299
93- if output_format == OutputFormats . JSON :
100+ if output_format == 'json' :
94101 return df .to_json (orient = orient )
95- elif output_format == OutputFormats . DATAFRAME :
102+ elif output_format == 'dataframe' :
96103 return df
97- elif output_format == OutputFormats . DICT :
104+ elif output_format == 'dict' :
98105 return df .to_dict ()
99106 raise TypeError ("Unknown output format" )
100107
@@ -117,7 +124,7 @@ def create_object(self, obj:object) -> Union[int, str]:
117124 return pkey_value
118125
119126 def read_object (self , model :Type [Base ], pkey :Union [int , str ],
120- output_format : OutputFormats = OutputFormats . DICT ,
127+ output_format :str = 'dict' ,
121128 orient :str = 'records' ) -> Dict [str , Any ]:
122129
123130 with self .session_maker () as session :
@@ -126,8 +133,14 @@ def read_object(self, model:Type[Base], pkey:Union[int, str],
126133 if obj is None : raise ObjectNotFound (f"No '{ model .__name__ } ' object found with { pkey_name } = { pkey } " )
127134 session .commit ()
128135
136+ # convert obj_dict to a dictionary if it isn't one already
129137 obj_dict = obj .to_dict ()
130- if output_format == OutputFormats .DICT :
138+
139+ # guard against invalid output_format strings
140+ if output_format not in OUTPUT_FORMATS :
141+ raise ValueError (f':param output_format = { output_format } , which is not one of the following valid output_format strings: { OUTPUT_FORMATS } ' )
142+
143+ if output_format == 'dict' :
131144 return obj_dict
132145
133146 else :
@@ -139,9 +152,9 @@ def read_object(self, model:Type[Base], pkey:Union[int, str],
139152 obj_dict [key ] = new_value
140153
141154 obj_df = pd .DataFrame .from_dict (obj_dict )
142- if output_format == OutputFormats . DATAFRAME :
155+ if output_format == 'dataframe' :
143156 return obj_df
144- elif output_format == OutputFormats . JSON :
157+ elif output_format == 'json' :
145158 return obj_df .to_json (orient = orient )
146159 raise TypeError ("Unknown output format" )
147160
0 commit comments