compress setup_query annotation

palewire · Feb 16, 2018 · 91f3a49 · 91f3a49
1 parent b3b45a0
commit 91f3a49
Showing 1 changed file with 5 additions and 4 deletions.
diff --git 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):