From 3b10879898adbfb6a1031777b339972f3772baba Mon Sep 17 00:00:00 2001
From: Matthew Chang <matt@realityengines.ai>
Date: Sat, 20 Feb 2021 03:55:31 +0000
Subject: [PATCH] Introduce new setup command to consolidate artifacts to a
 namespace

Introduce a new `restructure` command that mangles the repo. This moves
the protobufs to directories inside the min_tfs_client namespace and
then does a string replace on all the protobuf imports to reflect the
new top level namespace.

This allows the package to provide the same, lightweight python runtime
as it was originally intended, while still being used in an image with
tensorflow installed.
---
 .gitignore |  3 +++
 README.md  |  4 +++-
 setup.py   | 35 +++++++++++++++++++++++++++++++++--
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 33ebe43685d..97d7616ee37 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,3 +13,6 @@ dist/
 # Minified Tensorflow / Tensor Serving Folders
 tensor_serving_client/tensorflow
 tensor_serving_client/tensorflow_serving
+tensor_serving_client/min_tfs_client/tensorflow
+tensor_serving_client/min_tfs_client/tensorflow_serving
+
diff --git a/README.md b/README.md
index 3c5e72d3502..9a11cfb29e0 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,6 @@
-![repo-checks](https://github.com/zendesk/min-tfs-client/workflows/repo-checks/badge.svg)
+A modified version of min-tf-client where it is in its own isolated namespace, allowing for use alongside a Tensorflow install.
+Introduces a new setup comand: `python setup.py restructure`, which mangles the imports and moves the tensorflow and tensorflow_serving protobufs inside the min_tfs_client namespace.
+
 # Minimal Tensor Serving Python Client
 A lightweight python client to communicate with Tensor Serving.
 
diff --git a/setup.py b/setup.py
index 7e81c1d41e3..14a9b577b01 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 from contextlib import contextmanager
 from distutils.cmd import Command
 from pathlib import Path
-from shutil import copy2
+from shutil import copy2, rmtree
 from subprocess import check_output
 from typing import Any, List
 
@@ -73,12 +73,42 @@ def run(self):
         copy2(model_service_stub, destination)
 
 
+class RestructureProtos(Command):
+    description = '''
+    Rearrange the repository and modify protobuf imports so that protobufs nested inside min_tfs_client namespace
+    '''
+    user_options: List[Any] = []
+
+    def initialize_options(self):
+        ...
+
+    def finalize_options(self):
+        ...
+
+    def run(self):
+        with cd(str(OUTPUT_PATH)):
+            # Move the protos inside of the min_tfs_client namespace
+            if os.path.isdir('min_tfs_client/tensorflow'):
+                rmtree('min_tfs_client/tensorflow')
+            os.rename('tensorflow', 'min_tfs_client/tensorflow')
+            if os.path.isdir('min_tfs_client/tensorflow_serving'):
+                rmtree('min_tfs_client/tensorflow_serving')
+            os.rename('tensorflow_serving', 'min_tfs_client/tensorflow_serving')
+        for file_path in OUTPUT_PATH.rglob('*.py'):
+            filename = str(file_path)
+            with open(filename, 'r') as f:
+                new_text = f.read().replace('from tensorflow', 'from min_tfs_client.tensorflow')
+            with open(filename, 'w') as f:
+                f.write(new_text)
+
+
 class BuildPyCommand(build_py):
     """Custom build command."""
 
     def run(self):
         self.run_command("compile_pb")
         self.run_command("copy_grpc")
+        self.run_command("restructure")
         build_py.run(self)
 
 
@@ -88,13 +118,14 @@ def run(self):
 
 setup(
     name="min_tfs_client",
-    version="1.0.2",
+    version="1.0.2-abacus",
     description="A minified Tensor Serving Client for Python",
     long_description=long_description,
     long_description_content_type="text/markdown",
     cmdclass={
         "compile_pb": CompileProtobufs,
         "copy_grpc": CopyGRPCStubs,
+        "restructure": RestructureProtos,
         "build_py": BuildPyCommand,
     },
     package_dir={"": "tensor_serving_client"},