pytorch

1
# Owner(s): ["module: unknown"]
2

3
import glob
4
import io
5
import os
6
import unittest
7

8
import torch
9
from torch.testing._internal.common_utils import TestCase, run_tests
10

11

12
try:
13
    from third_party.build_bundled import create_bundled
14
except ImportError:
15
    create_bundled = None
16

17
license_file = 'third_party/LICENSES_BUNDLED.txt'
18
starting_txt = 'The Pytorch repository and source distributions bundle'
19
site_packages = os.path.dirname(os.path.dirname(torch.__file__))
20
distinfo = glob.glob(os.path.join(site_packages, 'torch-*dist-info'))
21

22
class TestLicense(TestCase):
23

24
    @unittest.skipIf(not create_bundled, "can only be run in a source tree")
25
    def test_license_for_wheel(self):
26
        current = io.StringIO()
27
        create_bundled('third_party', current)
28
        with open(license_file) as fid:
29
            src_tree = fid.read()
30
        if not src_tree == current.getvalue():
31
            raise AssertionError(
32
                f'the contents of "{license_file}" do not '
33
                'match the current state of the third_party files. Use '
34
                '"python third_party/build_bundled.py" to regenerate it')
35

36
    @unittest.skipIf(len(distinfo) == 0, "no installation in site-package to test")
37
    def test_distinfo_license(self):
38
        """If run when pytorch is installed via a wheel, the license will be in
39
        site-package/torch-*dist-info/LICENSE. Make sure it contains the third
40
        party bundle of licenses"""
41

42
        if len(distinfo) > 1:
43
            raise AssertionError('Found too many "torch-*dist-info" directories '
44
                                 f'in "{site_packages}, expected only one')
45
        with open(os.path.join(os.path.join(distinfo[0], 'LICENSE'))) as fid:
46
            txt = fid.read()
47
            self.assertTrue(starting_txt in txt)
48

49
if __name__ == '__main__':
50
    run_tests()
51