diff --git a/setup.py b/setup.py
index a2ca3137..9990f7de 100644
--- a/setup.py
+++ b/setup.py
@@ -9,11 +9,12 @@
     author_email="awohns@gmail.com",
     description="Infer node ages from a tree sequence topology",
     long_description=long_description,
+    packages=["tsdate"],
     long_description_content_type="text/markdown",
     url="http://pypi.python.org/pypi/tsdate",
     entry_points={
         'console_scripts': [
-            'tsdate=tsdate.cli:tsdate_main',
+            'tsdate=tsdate.__main__:main',
         ]
     },
     setup_requires=['setuptools_scm'],
diff --git a/tsdate/__init__.py b/tsdate/__init__.py
index e9fd567f..df67f920 100644
--- a/tsdate/__init__.py
+++ b/tsdate/__init__.py
@@ -20,9 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from .date import (  # NOQA
-    date, build_prior_grid
-    )
 
+from .date import date, build_prior_grid # NOQA
 
-from .provenance import __version__  # NOQA
+from .provenance import __version__ # NOQA
diff --git a/tsdate/__main__.py b/tsdate/__main__.py
index 0fd0e628..9261ffd3 100644
--- a/tsdate/__main__.py
+++ b/tsdate/__main__.py
@@ -1,4 +1,4 @@
-import tsdate.cli as cli
+from . import cli
 
 
 def main():