diff --git a/postgres_copy/copy_to.py b/postgres_copy/copy_to.py index fb75872..cc7be1e 100644 --- a/postgres_copy/copy_to.py +++ b/postgres_copy/copy_to.py @@ -27,10 +27,11 @@ def setup_query(self): for field in self.query.copy_to_fields: # raises error if field is not available expression = self.query.resolve_ref(field) - if field in self.query.annotations: - selection = (expression, self.compile(expression), field) - else: - selection = (expression, self.compile(expression), None) + selection = ( + expression, + self.compile(expression), + field if field in self.query.annotations else None, + ) self.select.append(selection) def execute_sql(self, csv_path=None):