diff --git a/swagger_marshmallow_codegen/codegen/config.py b/swagger_marshmallow_codegen/codegen/config.py
index 429da80..56f8b2f 100644
--- a/swagger_marshmallow_codegen/codegen/config.py
+++ b/swagger_marshmallow_codegen/codegen/config.py
@@ -2,10 +2,11 @@ from __future__ import annotations
 import typing_extensions as tx
 
 
-class ConfigDict(tx.TypedDict):
+class ConfigDict(tx.TypedDict, total=False):
     schema: bool
     input: bool
     output: bool
 
     emit_schema_even_primitive_type: bool
     skip_header_comment: bool
+    header_comment: str
diff --git a/swagger_marshmallow_codegen/codegen/v2/codegen.py b/swagger_marshmallow_codegen/codegen/v2/codegen.py
index 568f432..ecf29f2 100644
--- a/swagger_marshmallow_codegen/codegen/v2/codegen.py
+++ b/swagger_marshmallow_codegen/codegen/v2/codegen.py
@@ -95,7 +95,9 @@ class SchemaWriter:
             logger.debug("      nested: %s, %s", caller_name, field_class_name)
             if opts:
                 kwargs = LazyFormat(", {}", kwargs)
-            value = LazyFormat("{}(lambda: {}(){})", caller_name, field_class_name, kwargs)
+            value = LazyFormat(
+                "{}(lambda: {}(){})", caller_name, field_class_name, kwargs
+            )
         else:
             if caller_name == "fields.Nested":
                 caller_name = "fields.Field"
@@ -488,9 +490,14 @@ class Codegen:
     def resolver(self) -> Resolver:
         return self.accessor.resolver
 
-    def write_header(self, c):
-        c.im.stmt("# -*- coding:utf-8 -*-")
-        c.im.stmt("# this is auto-generated by swagger-marshmallow-codegen")
+    def write_header(self, c, *, comment: t.Optional[str] = None):
+        if comment is None:
+            comment = """\
+# this is auto-generated by swagger-marshmallow-codegen
+from __future__ import annotations
+"""
+        for line in comment.splitlines():
+            c.im.stmt(line)
 
     def write_import_(self, c):
         c.from_(*self.schema_class_path.rsplit(":", 1))
@@ -509,7 +516,7 @@ class Codegen:
     def codegen(self, d, ctx=None):
         c = ctx or Context()
         if not self.accessor.config.get("skip_header_comment", False):
-            self.write_header(c)
+            self.write_header(c, comment=self.accessor.config.get("header_comment"))
         c.m.sep()
         self.write_import_(c)
         self.write_body(c, d)
diff --git a/swagger_marshmallow_codegen/codegen/v3/codegen.py b/swagger_marshmallow_codegen/codegen/v3/codegen.py
index 4b8d214..0e27d70 100644
--- a/swagger_marshmallow_codegen/codegen/v3/codegen.py
+++ b/swagger_marshmallow_codegen/codegen/v3/codegen.py
@@ -1,1 +1,2 @@
 from ..v2.codegen import Codegen
+__all__ = ["Codegen"]
diff --git a/swagger_marshmallow_codegen/dispatcher.py b/swagger_marshmallow_codegen/dispatcher.py
index a5252be..2bf21f1 100644
--- a/swagger_marshmallow_codegen/dispatcher.py
+++ b/swagger_marshmallow_codegen/dispatcher.py
@@ -22,7 +22,7 @@ TYPE_MAP = {
     Pair(type="string", format=None): "marshmallow.fields:String",
     Pair(type="boolean", format=None): "marshmallow.fields:Boolean",
     Pair(type="string", format="uuid"): "marshmallow.fields:UUID",
-    Pair(type="string", format="date-time"): "marshmallow.fields:DateTime",
+    Pair(type="string", format="date-time"): "marshmallow.fields:AwareDateTime",
     Pair(type="string", format="date"): "marshmallow.fields:Date",
     Pair(type="string", format="time"): "marshmallow.fields:Time",
     Pair(type="string", format="email"): "marshmallow.fields:Email",
diff --git a/swagger_marshmallow_codegen/resolver.py b/swagger_marshmallow_codegen/resolver.py
index 844e9a2..50a350c 100644
--- a/swagger_marshmallow_codegen/resolver.py
+++ b/swagger_marshmallow_codegen/resolver.py
@@ -1,6 +1,5 @@
 # -*- coding:utf-8 -*-
 import logging
-import sys
 from collections import OrderedDict
 import dictknife
 from .langhelpers import titleize, normalize

diff --git a/swagger_marshmallow_codegen/tests/legacy_dst/00default.py b/swagger_marshmallow_codegen/tests/legacy_dst/00default.py
index fc76808..090bfa1 100644
--- a/swagger_marshmallow_codegen/tests/legacy_dst/00default.py
+++ b/swagger_marshmallow_codegen/tests/legacy_dst/00default.py
@@ -10,7 +10,8 @@ class X(Schema):
     string = fields.String(missing=lambda: 'default')
     integer = fields.Integer(missing=lambda: 10)
     boolean = fields.Boolean(missing=lambda: True)
-    datetime = fields.DateTime(missing=lambda: datetime.datetime(2000, 1, 1, 1, 1, 1))
+    datetime = fields.AwareDateTime(missing=lambda: datetime.datetime(2000, 1, 1, 1, 1, 1, tzinfo=datetime.timezone.utc))
+    date = fields.Date(missing=lambda: datetime.date(2000, 1, 1))
     object = fields.Nested(lambda: XObject(), missing=lambda: OrderedDict([('name', 'foo'), ('age', 20)]))
     array = fields.List(fields.Integer(), missing=lambda: [1, 2, 3])
 
diff --git a/swagger_marshmallow_codegen/tests/legacy_dst/00empty.py b/swagger_marshmallow_codegen/tests/legacy_dst/00empty.py
index c71e243..0890aac 100644
--- a/swagger_marshmallow_codegen/tests/legacy_dst/00empty.py
+++ b/swagger_marshmallow_codegen/tests/legacy_dst/00empty.py
@@ -1,3 +1,4 @@
+# flake8: noqa
 from marshmallow import (
     Schema,
     fields,
diff --git a/swagger_marshmallow_codegen/tests/legacy_dst/00paths.py b/swagger_marshmallow_codegen/tests/legacy_dst/00paths.py
index b8effe3..3b4e48e 100644
--- a/swagger_marshmallow_codegen/tests/legacy_dst/00paths.py
+++ b/swagger_marshmallow_codegen/tests/legacy_dst/00paths.py
@@ -15,7 +15,7 @@ class Pet(Schema):
     name = fields.String(required=True, description="Pet's name", validate=[Length(min=1, max=100, equal=None)])
     animal_type = fields.String(required=True, description='Kind of animal', validate=[Length(min=1, max=None, equal=None)])
     tags = fields.Field(description='Custom tags')
-    created = fields.DateTime(description='Creation time', dump_only=True)
+    created = fields.AwareDateTime(description='Creation time', dump_only=True)
 
 
 class PetsInput:
diff --git a/swagger_marshmallow_codegen/tests/legacy_src/00default.yaml b/swagger_marshmallow_codegen/tests/legacy_src/00default.yaml
index df6f0e9..cad6cef 100644
--- a/swagger_marshmallow_codegen/tests/legacy_src/00default.yaml
+++ b/swagger_marshmallow_codegen/tests/legacy_src/00default.yaml
@@ -14,6 +14,10 @@ definitions:
         type: string
         format: date-time
         default: 2000-01-01T01:01:01Z
+      date:
+        type: string
+        format: date
+        default: 2000-01-01
       object:
         type: object
         properties:
diff --git a/swagger_marshmallow_codegen/tests/test_codegen_legacy.py b/swagger_marshmallow_codegen/tests/test_codegen_legacy.py
index b3f67b9..d87dbd0 100644
--- a/swagger_marshmallow_codegen/tests/test_codegen_legacy.py
+++ b/swagger_marshmallow_codegen/tests/test_codegen_legacy.py
@@ -8,51 +8,53 @@ here = pathlib.Path(__file__).parent
 
 
 @pytest.mark.parametrize(
-    "src_file, dst_file",
+    "src_file, dst_file, header_comment",
     [
-        ("./legacy_src/00person.yaml", "./legacy_dst/00person.py"),
-        ("./legacy_src/01person.yaml", "./legacy_dst/01person.py"),
-        ("./legacy_src/02person.yaml", "./legacy_dst/02person.py"),
-        ("./legacy_src/03person.yaml", "./legacy_dst/03person.py"),
-        ("./legacy_src/04person.yaml", "./legacy_dst/04person.py"),
-        ("./legacy_src/05person.yaml", "./legacy_dst/05person.py"),
-        ("./legacy_src/00commit.yaml", "./legacy_dst/00commit.py"),
-        ("./legacy_src/01commit.yaml", "./legacy_dst/01commit.py"),
-        ("./legacy_src/00emojis.yaml", "./legacy_dst/00emojis.py"),
-        ("./legacy_src/00stat.yaml", "./legacy_dst/00stat.py"),
-        ("./legacy_src/00default.yaml", "./legacy_dst/00default.py"),
-        ("./legacy_src/00maximum.yaml", "./legacy_dst/00maximum.py"),
-        ("./legacy_src/00length.yaml", "./legacy_dst/00length.py"),
-        ("./legacy_src/00regex.yaml", "./legacy_dst/00regex.py"),
-        ("./legacy_src/00enum.yaml", "./legacy_dst/00enum.py"),
-        ("./legacy_src/00items.yaml", "./legacy_dst/00items.py"),
-        ("./legacy_src/00readonly.yaml", "./legacy_dst/00readonly.py"),
-        ("./legacy_src/00allOf.yaml", "./legacy_dst/00allOf.py"),
-        ("./legacy_src/00allOf2.yaml", "./legacy_dst/00allOf2.py"),
-        ("./legacy_src/01allOf2.yaml", "./legacy_dst/01allOf2.py"),
-        ("./legacy_src/02allOf2.yaml", "./legacy_dst/02allOf2.py"),
-        ("./legacy_src/00paths.yaml", "./legacy_dst/00paths.py"),
-        ("./legacy_src/01paths.yaml", "./legacy_dst/01paths.py"),
-        ("./legacy_src/02paths.yaml", "./legacy_dst/02paths.py"),
-        ("./legacy_src/03paths.yaml", "./legacy_dst/03paths.py"),
-        ("./legacy_src/00empty.yaml", "./legacy_dst/00empty.py"),
-        ("./legacy_src/01empty.yaml", "./legacy_dst/01empty.py"),
+        ("./legacy_src/00person.yaml", "./legacy_dst/00person.py", ""),
+        ("./legacy_src/01person.yaml", "./legacy_dst/01person.py", ""),
+        ("./legacy_src/02person.yaml", "./legacy_dst/02person.py", ""),
+        ("./legacy_src/03person.yaml", "./legacy_dst/03person.py", ""),
+        ("./legacy_src/04person.yaml", "./legacy_dst/04person.py", ""),
+        ("./legacy_src/05person.yaml", "./legacy_dst/05person.py", ""),
+        ("./legacy_src/00commit.yaml", "./legacy_dst/00commit.py", ""),
+        ("./legacy_src/01commit.yaml", "./legacy_dst/01commit.py", ""),
+        ("./legacy_src/00emojis.yaml", "./legacy_dst/00emojis.py", ""),
+        ("./legacy_src/00stat.yaml", "./legacy_dst/00stat.py", ""),
+        ("./legacy_src/00default.yaml", "./legacy_dst/00default.py", ""),
+        ("./legacy_src/00maximum.yaml", "./legacy_dst/00maximum.py", ""),
+        ("./legacy_src/00length.yaml", "./legacy_dst/00length.py", ""),
+        ("./legacy_src/00regex.yaml", "./legacy_dst/00regex.py", ""),
+        ("./legacy_src/00enum.yaml", "./legacy_dst/00enum.py", ""),
+        ("./legacy_src/00items.yaml", "./legacy_dst/00items.py", ""),
+        ("./legacy_src/00readonly.yaml", "./legacy_dst/00readonly.py", ""),
+        ("./legacy_src/00allOf.yaml", "./legacy_dst/00allOf.py", ""),
+        ("./legacy_src/00allOf2.yaml", "./legacy_dst/00allOf2.py", ""),
+        ("./legacy_src/01allOf2.yaml", "./legacy_dst/01allOf2.py", ""),
+        ("./legacy_src/02allOf2.yaml", "./legacy_dst/02allOf2.py", ""),
+        ("./legacy_src/00paths.yaml", "./legacy_dst/00paths.py", ""),
+        ("./legacy_src/01paths.yaml", "./legacy_dst/01paths.py", ""),
+        ("./legacy_src/02paths.yaml", "./legacy_dst/02paths.py", ""),
+        ("./legacy_src/03paths.yaml", "./legacy_dst/03paths.py", ""),
+        ("./legacy_src/00empty.yaml", "./legacy_dst/00empty.py", "# flake8: noqa"),
+        ("./legacy_src/01empty.yaml", "./legacy_dst/01empty.py", ""),
         (
             "./legacy_src/00list_with_options.yaml",
             "./legacy_dst/00list_with_options.py",
+            "",
         ),
-        ("./legacy_src/00reserved.yaml", "./legacy_dst/00reserved.py"),
-        ("./legacy_src/00typearray.yaml", "./legacy_dst/00typearray.py"),
-        ("./legacy_src/00additional.yaml", "./legacy_dst/00additional.py"),
-        ("./legacy_src/01additional.yaml", "./legacy_dst/01additional.py"),
-        ("./legacy_src/00nullable.yaml", "./legacy_dst/00nullable.py"),
-        ("./legacy_src/00primitiveapi.yaml", "./legacy_dst/00primitiveapi.py"),
+        ("./legacy_src/00reserved.yaml", "./legacy_dst/00reserved.py", ""),
+        ("./legacy_src/00typearray.yaml", "./legacy_dst/00typearray.py", ""),
+        ("./legacy_src/00additional.yaml", "./legacy_dst/00additional.py", ""),
+        ("./legacy_src/01additional.yaml", "./legacy_dst/01additional.py", ""),
+        ("./legacy_src/00nullable.yaml", "./legacy_dst/00nullable.py", ""),
+        ("./legacy_src/00primitiveapi.yaml", "./legacy_dst/00primitiveapi.py", ""),
         # ("./legacy_src/00patternProperties.yaml", "./legacy_dst/00patternProperties.py"),  not supported yet
     ],
 )
-def test_v2(
-    src_file,
-    dst_file,
+def test(
+    src_file: str,
+    dst_file: str,
+    header_comment: str,
 ):
     from swagger_marshmallow_codegen.lifting import lifting_definition
     from swagger_marshmallow_codegen.codegen import Context
@@ -62,9 +64,8 @@ def test_v2(
 
     get_codegen().codegen(
         lifting_definition(d),
-        {"schema": True, "input": True, "output": True},
+        {"schema": True, "input": True, "output": True, "header_comment": header_comment},
         ctx=ctx,
-        test=True,
     )
 
     expected = load_dstfile(dst_file, here=here).rstrip("\n")

CI is broken, DateTime field handling is changed
refs https://github.com/marshmallow-code/marshmallow/issues/1234
swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00person.yaml-./legacy_dst/00person.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/01person.yaml-./legacy_dst/01person.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/02person.yaml-./legacy_dst/02person.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/03person.yaml-./legacy_dst/03person.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/04person.yaml-./legacy_dst/04person.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/05person.yaml-./legacy_dst/05person.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00commit.yaml-./legacy_dst/00commit.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/01commit.yaml-./legacy_dst/01commit.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00emojis.yaml-./legacy_dst/00emojis.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00stat.yaml-./legacy_dst/00stat.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00maximum.yaml-./legacy_dst/00maximum.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00length.yaml-./legacy_dst/00length.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00regex.yaml-./legacy_dst/00regex.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00enum.yaml-./legacy_dst/00enum.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00items.yaml-./legacy_dst/00items.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00readonly.yaml-./legacy_dst/00readonly.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00allOf.yaml-./legacy_dst/00allOf.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00allOf2.yaml-./legacy_dst/00allOf2.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/01allOf2.yaml-./legacy_dst/01allOf2.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/02allOf2.yaml-./legacy_dst/02allOf2.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00paths.yaml-./legacy_dst/00paths.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/01paths.yaml-./legacy_dst/01paths.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/02paths.yaml-./legacy_dst/02paths.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/03paths.yaml-./legacy_dst/03paths.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00empty.yaml-./legacy_dst/00empty.py-# swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/01empty.yaml-./legacy_dst/01empty.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00list_with_options.yaml-./legacy_dst/00list_with_options.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00reserved.yaml-./legacy_dst/00reserved.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00typearray.yaml-./legacy_dst/00typearray.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00additional.yaml-./legacy_dst/00additional.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/01additional.yaml-./legacy_dst/01additional.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00nullable.yaml-./legacy_dst/00nullable.py-] swagger_marshmallow_codegen/tests/test_codegen_legacy.py::test[./legacy_src/00primitiveapi.yaml-./legacy_dst/00primitiveapi.py-]