diff --git a/bloodhound_import/__init__.py b/bloodhound_import/__init__.py
index 96eff44..2472310 100644
--- a/bloodhound_import/__init__.py
+++ b/bloodhound_import/__init__.py
@@ -1,4 +1,5 @@
 import argparse
+import asyncio
 import logging
 from bloodhound_import import database
 from bloodhound_import.importer import parse_file, add_constraints
@@ -51,4 +52,4 @@ async def main():
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())
diff --git a/bloodhound_import/database.py b/bloodhound_import/database.py
index 7a994b6..d0a303d 100644
--- a/bloodhound_import/database.py
+++ b/bloodhound_import/database.py
@@ -2,13 +2,13 @@
 Code borrowed from
 https://github.com/fox-it/aclpwn.py/blob/master/aclpwn/database.py
 """
-from neo4j import AsyncGraphDatabase
+from neo4j import AsyncGraphDatabase, AsyncDriver
 import platform
 import json
 import os
 
 
-def init_driver(ip, port, scheme, user, password):
+def init_driver(ip, port, scheme, user, password) -> AsyncDriver:
     uri = "{}://{}:{}".format(scheme, ip, port)
     driver = AsyncGraphDatabase.driver(uri, auth=(user, password))
     return driver
diff --git a/bloodhound_import/importer.py b/bloodhound_import/importer.py
index 73d9194..b5a23a4 100644
--- a/bloodhound_import/importer.py
+++ b/bloodhound_import/importer.py
@@ -343,7 +343,7 @@ async def parse_container(tx: neo4j.Transaction, container: dict):
             query = build_add_edge_query('Container', target['ObjectType'], 'Contains', '{isacl: false}')
             await tx.run(query, props=dict(source=identifier, target=target['ObjectIdentifier']))
 
-async def parse_file(filename: str, driver: neo4j.GraphDatabase):
+async def parse_file(filename: str, driver: neo4j.AsyncDriver):
     """Parse a bloodhound file.
 
     Arguments:
diff --git a/setup.py b/setup.py
index f647f6d..4814363 100644
--- a/setup.py
+++ b/setup.py
@@ -16,6 +16,6 @@
     packages=['bloodhound_import'],
     install_requires=['neo4j-driver', 'ijson'],
     entry_points={
-        'console_scripts': ['bloodhound-import=bloodhound_import:main']
+        'console_scripts': ['bloodhound-import=bloodhound_import:__main__']
     }
 )