61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
from mongoengine.fields import DateField
|
|
from datetime import datetime
|
|
|
|
|
|
DEFAULT_DATE_FORMAT = '%b %d, %Y'
|
|
|
|
|
|
class QuarterDateField(DateField):
|
|
def __init__(self, *args, date_format='%b %d, %Y', **kwargs):
|
|
self.date_format = date_format
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def validate(self, value):
|
|
if isinstance(value, str):
|
|
try:
|
|
datetime.strptime(value, self.date_format)
|
|
except Exception as ex:
|
|
self.error(f'Date time is not the correct format: (expected {self.date_format}). Error: {str(ex)}')
|
|
|
|
def to_mongo(self, value):
|
|
value = super().to_mongo(value)
|
|
if isinstance(value, datetime):
|
|
value = value.strftime(self.date_format)
|
|
return value
|
|
|
|
def to_python(self, value):
|
|
value = super().to_python(value)
|
|
if isinstance(value, datetime):
|
|
value = value.strftime(self.date_format)
|
|
return value
|
|
|
|
def get_dt_format(self):
|
|
return self.date_format
|
|
|
|
@staticmethod
|
|
def make_datetime(value, dt_format='%b %d, %Y'):
|
|
if isinstance(value, str):
|
|
return datetime.strptime(value, dt_format)
|
|
return value
|
|
|
|
|
|
class QuarterYearField(QuarterDateField):
|
|
def __init__(self, *args, **kwargs):
|
|
kwargs['date_format'] = '%Y'
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def validate(self, value):
|
|
if isinstance(value, int):
|
|
value = str(value)
|
|
return super().validate(value)
|
|
|
|
def to_mongo(self, value):
|
|
if isinstance(value, int):
|
|
value = str(value)
|
|
return super().to_mongo(value)
|
|
|
|
def to_python(self, value):
|
|
if isinstance(value, int):
|
|
value = str(value)
|
|
return super().to_python(value)
|