1#!/pxrpythonsubst
2#
3# Copyright 2018 Pixar
4#
5# Licensed under the Apache License, Version 2.0 (the "Apache License")
6# with the following modification; you may not use this file except in
7# compliance with the Apache License and the following modification to it:
8# Section 6. Trademarks. is deleted and replaced with:
9#
10# 6. Trademarks. This License does not grant permission to use the trade
11#    names, trademarks, service marks, or product names of the Licensor
12#    and its affiliates, except as required to comply with Section 4(c) of
13#    the License and to reproduce the content of the NOTICE file.
14#
15# You may obtain a copy of the Apache License at
16#
17#     http://www.apache.org/licenses/LICENSE-2.0
18#
19# Unless required by applicable law or agreed to in writing, software
20# distributed under the Apache License with the above modification is
21# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
22# KIND, either express or implied. See the Apache License for the specific
23# language governing permissions and limitations under the Apache License.
24#
25from __future__ import print_function
26import argparse
27import glob
28import os
29import sys
30
31from pxr import Ar, Sdf, Tf, Usd, UsdUtils
32from contextlib import contextmanager
33
34@contextmanager
35def _Stream(path, *args, **kwargs):
36    if path == '-':
37        yield sys.stdout
38    else:
39        with open(path, *args, **kwargs) as fp:
40            yield fp
41
42def _Print(stream, msg):
43    print(msg, file=stream)
44
45def _Err(msg):
46    sys.stderr.write(msg + '\n')
47
48def _CheckCompliance(rootLayer, arkit=False):
49    checker = UsdUtils.ComplianceChecker(arkit=arkit,
50        # We're going to flatten the USD stage and convert the root layer to
51        # crate file format before packaging, if necessary. Hence, skip these
52        # checks.
53        skipARKitRootLayerCheck=True)
54    checker.CheckCompliance(rootLayer)
55    errors = checker.GetErrors()
56    failedChecks = checker.GetFailedChecks()
57    warnings = checker.GetWarnings()
58    for msg in errors + failedChecks:
59        _Err(msg)
60    if len(warnings) > 0:
61        _Err("*********************************************\n"
62             "Possible correctness problems to investigate:\n"
63             "*********************************************\n")
64        for msg in warnings:
65            _Err(msg)
66    return len(errors) == 0 and len(failedChecks) == 0
67
68def _CreateUsdzPackage(usdzFile, filesToAdd, recurse, checkCompliance, verbose):
69    with Usd.ZipFileWriter.CreateNew(usdzFile) as usdzWriter:
70        fileList = []
71        while filesToAdd:
72            # Pop front (first file) from the list of files to add.
73            f = filesToAdd[0]
74            filesToAdd = filesToAdd[1:]
75
76            if os.path.isdir(f):
77                # If a directory is specified, add all files in the directory.
78                filesInDir = glob.glob(os.path.join(f, '*'))
79                # If the recurse flag is not specified, remove sub-directories.
80                if not recurse:
81                    filesInDir = [f for f in filesInDir if not os.path.isdir(f)]
82                # glob.glob returns files in arbitrary order. Hence, sort them
83                # here to get consistent ordering of files in the package.
84                filesInDir.sort()
85                filesToAdd += filesInDir
86            else:
87                if verbose:
88                    print('.. adding: %s' % f)
89                if os.path.getsize(f) > 0:
90                    fileList.append(f)
91                else:
92                    _Err("Skipping empty file '%s'." % f)
93
94        if checkCompliance and len(fileList) > 0:
95            rootLayer = fileList[0]
96            if not _CheckCompliance(rootLayer):
97                return False
98
99        for f in fileList:
100            try:
101                usdzWriter.AddFile(f)
102            except Tf.ErrorException as e:
103                _Err('Failed to add file \'%s\' to package. Discarding '
104                    'package.' % f)
105                # When the "with" block exits, Discard() will be called on
106                # usdzWriter automatically if an exception occurs.
107                raise
108        return True
109
110def _DumpContents(dumpLocation, zipFile):
111    with _Stream(dumpLocation, "w") as ofp:
112        _Print(ofp, "    Offset\t      Comp\t    Uncomp\tName")
113        _Print(ofp, "    ------\t      ----\t    ------\t----")
114        fileNames = zipFile.GetFileNames()
115        for fileName in fileNames:
116            fileInfo = zipFile.GetFileInfo(fileName)
117            _Print(ofp, "%10d\t%10d\t%10d\t%s" %
118                (fileInfo.dataOffset, fileInfo.size,
119                    fileInfo.uncompressedSize, fileName))
120
121        _Print(ofp, "----------")
122        _Print(ofp, "%d files total" % len(fileNames))
123
124def _ListContents(listLocation, zipFile):
125    with _Stream(listLocation, "w") as ofp:
126        for fileName in zipFile.GetFileNames():
127            _Print(ofp, fileName)
128
129def main():
130    parser = argparse.ArgumentParser(description='Utility for creating a .usdz '
131        'file containing USD assets and for inspecting existing .usdz files.')
132
133    parser.add_argument('usdzFile', type=str, nargs='?',
134                        help='Name of the .usdz file to create or to inspect '
135                        'the contents of.')
136
137    parser.add_argument('inputFiles', type=str, nargs='*',
138                        help='Files to include in the .usdz file.')
139    parser.add_argument('-r', '--recurse', dest='recurse', action='store_true',
140                        help='If specified, files in sub-directories are '
141                        'recursively added to the package.')
142
143    parser.add_argument('-a', '--asset', dest='asset', type=str,
144                        help='Resolvable asset path pointing to the root layer '
145                        'of the asset to be isolated and copied into the '
146                        'package.')
147    parser.add_argument("--arkitAsset", dest="arkitAsset", type=str,
148                        help="Similar to the --asset option, the --arkitAsset "
149                        "option packages all of the dependencies of the named "
150                        "scene file.  Assets targeted at the initial usdz "
151                        "implementation in ARKit operate under greater "
152                        "constraints than usdz files for more general 'in "
153                        "house' uses, and this option attempts to ensure that "
154                        "these constraints are honored; this may involve more "
155                        "transformations to the data, which may cause loss of "
156                        "features such as VariantSets.")
157
158    parser.add_argument('-c', '--checkCompliance', dest='checkCompliance',
159                        action='store_true', help='Perform compliance checking '
160                        'of the input files. If the input asset or \"root\" '
161                        'layer fails any of the compliance checks, the package '
162                        'is not created and the program fails.')
163
164    parser.add_argument('-l', '--list', dest='listTarget', type=str,
165                        nargs='?', default=None, const='-',
166                        help='List contents of the specified usdz file. If '
167                        'a file-path argument is provided, the list is output '
168                        'to a file at the given path. If no argument is '
169                        'provided or if \'-\' is specified as the argument, the'
170                        ' list is output to stdout.')
171    parser.add_argument('-d', '--dump', dest='dumpTarget', type=str,
172                        nargs='?', default=None, const='-',
173                        help='Dump contents of the specified usdz file. If '
174                        'a file-path argument is provided, the contents are '
175                        'output to a file at the given path. If no argument is '
176                        'provided or if \'-\' is specified as the argument, the'
177                        ' contents are output to stdout.')
178
179    parser.add_argument('-v', '--verbose', dest='verbose', action='store_true',
180                        help='Enable verbose mode, which causes messages '
181                        'regarding files being added to the package to be '
182                        'output to stdout.')
183
184    args = parser.parse_args()
185    usdzFile = args.usdzFile
186    inputFiles = args.inputFiles
187
188    if args.asset and args.arkitAsset:
189        parser.error("Specify either --asset or --arkitAsset, not both.")
190
191    elif (args.arkitAsset or args.asset) and len(inputFiles)>0:
192        parser.error("Specify either inputFiles or an asset (via --asset or "
193                     "--arkitAsset, not both.")
194
195    # If usdzFile is not specified directly as an argument, check if it has been
196    # specified as an argument to the --list or --dump options. In these cases,
197    # output the list or the contents to stdout.
198    if not usdzFile:
199        if args.listTarget and args.listTarget != '-' and \
200           args.listTarget.endswith('.usdz') and \
201           os.path.exists(args.listTarget):
202            usdzFile = args.listTarget
203            args.listTarget = '-'
204        elif args.dumpTarget and args.dumpTarget != '-' and \
205           args.dumpTarget.endswith('.usdz') and \
206           os.path.exists(args.dumpTarget):
207            usdzFile = args.dumpTarget
208            args.dumpTarget = '-'
209        else:
210            parser.error("No usdz file specified!")
211
212    # Check if we're in package creation mode and verbose mode is enabled,
213    # print some useful information.
214    if (args.asset or args.arkitAsset or len(inputFiles)>0):
215        # Ensure that the usdz file has the right extension.
216        if not usdzFile.endswith('.usdz'):
217            usdzFile += '.usdz'
218
219        if args.verbose:
220            if os.path.exists(usdzFile):
221                print("File at path '%s' already exists. Overwriting file." %
222                        usdzFile)
223
224            if args.inputFiles:
225                print('Creating package \'%s\' with files %s.' %
226                      (usdzFile, inputFiles))
227
228            if args.asset or args.arkitAsset:
229                Tf.Debug.SetDebugSymbolsByName("USDUTILS_CREATE_USDZ_PACKAGE", 1)
230
231            if not args.recurse:
232                print('Not recursing into sub-directories.')
233    else:
234        if args.checkCompliance:
235            parser.error("--checkCompliance should only be specified when "
236                "creating a usdz package. Please use 'usdchecker' to check "
237                "compliance of an existing .usdz file.")
238
239
240    success = True
241    if len(inputFiles) > 0:
242        success = _CreateUsdzPackage(usdzFile, inputFiles, args.recurse,
243                args.checkCompliance, args.verbose) and success
244
245    elif args.asset:
246        r = Ar.GetResolver()
247        resolvedAsset = r.Resolve(args.asset)
248        if args.checkCompliance:
249            success = _CheckCompliance(resolvedAsset, arkit=False) and success
250
251        context = r.CreateDefaultContextForAsset(resolvedAsset)
252        with Ar.ResolverContextBinder(context):
253            # Create the package only if the compliance check was passed.
254            success = success and UsdUtils.CreateNewUsdzPackage(
255                Sdf.AssetPath(args.asset), usdzFile)
256
257    elif args.arkitAsset:
258        r = Ar.GetResolver()
259        resolvedAsset = r.Resolve(args.arkitAsset)
260        if args.checkCompliance:
261            success = _CheckCompliance(resolvedAsset, arkit=True) and success
262
263        context = r.CreateDefaultContextForAsset(resolvedAsset)
264        with Ar.ResolverContextBinder(context):
265            # Create the package only if the compliance check was passed.
266            success = success and UsdUtils.CreateNewARKitUsdzPackage(
267                    Sdf.AssetPath(args.arkitAsset), usdzFile)
268
269    if args.listTarget or args.dumpTarget:
270        if os.path.exists(usdzFile):
271            zipFile = Usd.ZipFile.Open(usdzFile)
272            if zipFile:
273                if args.dumpTarget:
274                    if args.dumpTarget == usdzFile:
275                        _Err("The file into which to dump the contents of the "
276                             "usdz file '%s' must be different from the file "
277                             "itself." % usdzFile)
278                        return 1
279                    _DumpContents(args.dumpTarget, zipFile)
280                if args.listTarget:
281                    if args.listTarget == usdzFile:
282                        _Err("The file into which to list the contents of the "
283                             "usdz file '%s' must be different from the file "
284                             "itself." % usdzFile)
285                        return 1
286                    _ListContents(args.listTarget, zipFile)
287            else:
288                _Err("Failed to open usdz file at path '%s'." % usdzFile)
289        else:
290            _Err("Can't find usdz file at path '%s'." % usdzFile)
291
292    return 0 if success else 1
293
294if __name__ == '__main__':
295    sys.exit(main())
296