From 503fb68db74dcd316f4d11b2e516e5ea03ff5fa3 Mon Sep 17 00:00:00 2001
From: Joyce Er <huer@microsoft.com>
Date: Wed, 2 Jun 2021 05:12:24 +0000
Subject: [PATCH] Add PEP 396 __version__ attribute

---
 tb_plugin/setup.py                      | 20 ++++++++++++++++----
 tb_plugin/torch_tb_profiler/__init__.py |  2 ++
 tb_plugin/version.txt                   |  1 -
 3 files changed, 18 insertions(+), 5 deletions(-)
 delete mode 100644 tb_plugin/version.txt

diff --git a/tb_plugin/setup.py b/tb_plugin/setup.py
index 7fed527e8..8880098b4 100644
--- a/tb_plugin/setup.py
+++ b/tb_plugin/setup.py
@@ -6,14 +6,24 @@
 import setuptools
 import subprocess
 
-def get_version():
-    with open("version.txt", encoding="utf-8") as f:
-        version = f.read().strip()
+
+def read(rel_path):
+    here = os.path.abspath(os.path.dirname(__file__))
+    with open(os.path.join(here, rel_path)) as fp:
+        return fp.read()
+
+
+def get_version(rel_path):
+    for line in read(rel_path).splitlines():
+        if line.startswith("__version__"):
+            delim = '"' if '"' in line else "'"
+            version = line.split(delim)[1]
 
     if os.getenv('TORCH_TB_PROFILER_BUILD_VERSION'):
         version = os.getenv('TORCH_TB_PROFILER_BUILD_VERSION')
     return version
 
+
 INSTALL_REQUIRED = [
     "pandas >= 1.0.0",
     "tensorboard >= 1.15, !=2.1.0"
@@ -30,6 +40,7 @@ def get_version():
     "gs": ["google-cloud-storage"]
 }
 
+
 class build_fe(setuptools.Command):
     """Build the frontend"""
     description = "run yarn build on frontend directory"
@@ -50,9 +61,10 @@ def run(self):
         # restore the working directory
         os.chdir(cwd)
 
+
 setuptools.setup(
     name="torch_tb_profiler",
-    version=get_version(),
+    version=get_version(os.path.join('torch_tb_profiler', '__init__.py')),
     description="PyTorch Profiler TensorBoard Plugin",
     long_description="PyTorch Profiler TensorBoard Plugin : \
         https://github.com/pytorch/kineto/tree/master/tb_plugin",
diff --git a/tb_plugin/torch_tb_profiler/__init__.py b/tb_plugin/torch_tb_profiler/__init__.py
index 51b365adf..5dcc1a1e7 100644
--- a/tb_plugin/torch_tb_profiler/__init__.py
+++ b/tb_plugin/torch_tb_profiler/__init__.py
@@ -3,3 +3,5 @@
 # --------------------------------------------------------------------------
 
 # Entry point for Pytorch TensorBoard plugin package.
+
+__version__ = "0.1.0"
diff --git a/tb_plugin/version.txt b/tb_plugin/version.txt
deleted file mode 100644
index 6e8bf73aa..000000000
--- a/tb_plugin/version.txt
+++ /dev/null
@@ -1 +0,0 @@
-0.1.0