concantenate LICENSE files when building a wheel (pytorch#51634)

Summary:
Fixes pytorch#50695

I checked locally that the concatenated license file appears at `torch-<version>.dist-info/LICENSE` in the wheel.

Pull Request resolved: pytorch#51634

Reviewed By: zhangguanheng66

Differential Revision: D26225550

Pulled By: walterddr

fbshipit-source-id: 830c59fb7aea0eb50b99e295edddad9edab6ba3a

Author: mattip

setup.py

@@ -552,6 +552,50 @@ def load(filename):
             with open('compile_commands.json', 'w') as f:
                 f.write(new_contents)
 
+class concat_license_files():
+    """Merge LICENSE and LICENSES_BUNDLED.txt as a context manager
+
+    LICENSE is the main PyTorch license, LICENSES_BUNDLED.txt is auto-generated
+    from all the licenses found in ./third_party/. We concatenate them so there
+    is a single license file in the sdist and wheels with all of the necessary
+    licensing info.
+    """
+    def __init__(self):
+        self.f1 = 'LICENSE'
+        self.f2 = 'third_party/LICENSES_BUNDLED.txt'
+
+    def __enter__(self):
+        """Concatenate files"""
+        with open(self.f1, 'r') as f1:
+            self.bsd_text = f1.read()
+
+        with open(self.f1, 'a') as f1:
+            with open(self.f2, 'r') as f2:
+                self.bundled_text = f2.read()
+                f1.write('\n\n')
+                f1.write(self.bundled_text)
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        """Restore content of f1"""
+        with open(self.f1, 'w') as f:
+            f.write(self.bsd_text)
+
+
+try:
+    from wheel.bdist_wheel import bdist_wheel
+except ImportError:
+    # This is useful when wheel is not installed and bdist_wheel is not
+    # specified on the command line. If it _is_ specified, parsing the command
+    # line will fail before wheel_concatenate is needed
+    wheel_concatenate = None
+else:
+    # Need to create the proper LICENSE.txt for the wheel
+    class wheel_concatenate(bdist_wheel):
+        """ check submodules on sdist to prevent incomplete tarballs """
+        def run(self):
+            with concat_license_files():
+                super().run()
+
 
 class install(setuptools.command.install.install):
     def run(self):
@@ -724,6 +768,7 @@ def make_relative_rpath_args(path):
         'build_ext': build_ext,
         'clean': clean,
         'install': install,
+        'bdist_wheel': wheel_concatenate,
     }
 
     entry_points = {

test/run_test.py

@@ -108,6 +108,7 @@
     'test_fx_experimental',
     'test_functional_autograd_benchmark',
     'test_package',
+    'test_license',
     'distributed/pipeline/sync/skip/test_api',
     'distributed/pipeline/sync/skip/test_gpipe',
     'distributed/pipeline/sync/skip/test_inspect_skip_layout',

test/test_license.py

@@ -1,6 +1,9 @@
+import glob
 import io
+import os
 import unittest
 
+import torch
 from torch.testing._internal.common_utils import TestCase, run_tests
 
 
@@ -10,11 +13,14 @@
     create_bundled = None
 
 license_file = 'third_party/LICENSES_BUNDLED.txt'
+starting_txt = 'The Pytorch repository and source distributions bundle'
+site_packages = os.path.dirname(os.path.dirname(torch.__file__))
+distinfo = glob.glob(os.path.join(site_packages, 'torch-*dist-info'))
 
 class TestLicense(TestCase):
 
     @unittest.skipIf(not create_bundled, "can only be run in a source tree")
-    def test_license_in_wheel(self):
+    def test_license_for_wheel(self):
         current = io.StringIO()
         create_bundled('third_party', current)
         with open(license_file) as fid:
@@ -25,6 +31,18 @@ def test_license_in_wheel(self):
                 'match the current state of the third_party files. Use '
                 '"python third_party/build_bundled.py" to regenerate it')
 
+    @unittest.skipIf(len(distinfo) == 0, "no installation in site-package to test")
+    def test_distinfo_license(self):
+        """If run when pytorch is installed via a wheel, the license will be in
+        site-package/torch-*dist-info/LICENSE. Make sure it contains the third
+        party bundle of licenses"""
+
+        if len(distinfo) > 1:
+            raise AssertionError('Found too many "torch-*dist-info" directories '
+                                 f'in "{site_packages}, expected only one')
+        with open(os.path.join(os.path.join(distinfo[0], 'LICENSE'))) as fid:
+            txt = fid.read()
+            self.assertTrue(starting_txt in txt)
 
 if __name__ == '__main__':
     run_tests()