Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Epi export rewrite #911

Merged
merged 13 commits into from
Oct 12, 2023
23 changes: 12 additions & 11 deletions hawc/apps/common/exports.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import pandas as pd
from django.db.models import QuerySet
from django.utils import timezone

from .helper import FlatExport

Expand All @@ -14,18 +15,10 @@ def __init__(
include: tuple[str, ...] | None = None,
exclude: tuple[str, ...] | None = None,
):
"""Instantiate an exporter instance for a given django model.

Args:
key_prefix (str, optional): The model name to prepend to data frame columns.
query_prefix (str, optional): The model prefix in the ORM.
include (tuple | None, optional): If included, only these items are added.
exclude (tuple | None, optional): If specified, items are removed from base.
"""
self.key_prefix = key_prefix + "-" if key_prefix else key_prefix
self.query_prefix = query_prefix + "__" if query_prefix else query_prefix
self.include = (key_prefix + field for field in include) if include else tuple()
self.exclude = (key_prefix + field for field in exclude) if exclude else tuple()
self.include = tuple(self.key_prefix + field for field in include) if include else tuple()
self.exclude = tuple(self.key_prefix + field for field in exclude) if exclude else tuple()

@property
def value_map(self) -> dict:
Expand Down Expand Up @@ -153,6 +146,15 @@ def prepare_df(self, df: pd.DataFrame) -> pd.DataFrame:
"""
return df

def format_time(self, df: pd.DataFrame) -> pd.DataFrame:
if df.shape[0] == 0:
return df
tz = timezone.get_default_timezone()
for key in [self.get_column_name("created"), self.get_column_name("last_updated")]:
if key in df.columns:
df.loc[:, key] = df[key].dt.tz_convert(tz).dt.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
return df

def get_df(self, qs: QuerySet) -> pd.DataFrame:
"""Get dataframe export from queryset.

Expand Down Expand Up @@ -211,7 +213,6 @@ def get_df(self, qs: QuerySet) -> pd.DataFrame:
@classmethod
def flat_export(cls, qs: QuerySet, filename: str) -> FlatExport:
"""Return an instance of a FlatExport.

Args:
qs (QuerySet): the initial QuerySet
filename (str): the filename for the export
Expand Down
Loading