17757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch"""Code unit (module) handling for Coverage."""
27757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
37757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdochimport glob, os
47757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
57757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdochfrom coverage.backward import open_source, string_class, StringIO
67757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdochfrom coverage.misc import CoverageException
77757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
87757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
97757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdochdef code_unit_factory(morfs, file_locator):
107757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    """Construct a list of CodeUnits from polymorphic inputs.
117757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
127757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    `morfs` is a module or a filename, or a list of same.
137757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
147757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    `file_locator` is a FileLocator that can help resolve filenames.
157757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
167757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    Returns a list of CodeUnit objects.
177757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
187757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    """
197757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    # Be sure we have a list.
207757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    if not isinstance(morfs, (list, tuple)):
217757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        morfs = [morfs]
227757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
237757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    # On Windows, the shell doesn't expand wildcards.  Do it here.
247757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    globbed = []
257757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    for morf in morfs:
267757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        if isinstance(morf, string_class) and ('?' in morf or '*' in morf):
277757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            globbed.extend(glob.glob(morf))
287757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        else:
297757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            globbed.append(morf)
307757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    morfs = globbed
317757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
327757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    code_units = [CodeUnit(morf, file_locator) for morf in morfs]
337757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
347757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    return code_units
357757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
367757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
377757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdochclass CodeUnit(object):
387757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    """Code unit: a filename or module.
397757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
407757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    Instance attributes:
417757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
427757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    `name` is a human-readable name for this code unit.
437757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    `filename` is the os path from which we can read the source.
447757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    `relative` is a boolean.
457757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
467757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    """
477757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def __init__(self, morf, file_locator):
487757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        self.file_locator = file_locator
497757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
507757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        if hasattr(morf, '__file__'):
517757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            f = morf.__file__
527757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        else:
537757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            f = morf
547757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        # .pyc files should always refer to a .py instead.
557757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        if f.endswith('.pyc'):
567757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            f = f[:-1]
577757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        self.filename = self.file_locator.canonical_filename(f)
587757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
597757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        if hasattr(morf, '__name__'):
607757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            n = modname = morf.__name__
617757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            self.relative = True
627757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        else:
637757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            n = os.path.splitext(morf)[0]
647757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            rel = self.file_locator.relative_filename(n)
657757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            if os.path.isabs(n):
667757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch                self.relative = (rel != n)
677757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            else:
687757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch                self.relative = True
697757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            n = rel
707757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            modname = None
717757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        self.name = n
727757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        self.modname = modname
737757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
747757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def __repr__(self):
757757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        return "<CodeUnit name=%r filename=%r>" % (self.name, self.filename)
767757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
777757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    # Annoying comparison operators. Py3k wants __lt__ etc, and Py2k needs all
787757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    # of them defined.
797757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
807757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def __lt__(self, other): return self.name <  other.name
817757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def __le__(self, other): return self.name <= other.name
827757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def __eq__(self, other): return self.name == other.name
837757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def __ne__(self, other): return self.name != other.name
847757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def __gt__(self, other): return self.name >  other.name
857757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def __ge__(self, other): return self.name >= other.name
867757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
877757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def flat_rootname(self):
887757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        """A base for a flat filename to correspond to this code unit.
897757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
907757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        Useful for writing files about the code where you want all the files in
917757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        the same directory, but need to differentiate same-named files from
927757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        different directories.
937757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
947757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        For example, the file a/b/c.py might return 'a_b_c'
957757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
967757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        """
977757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        if self.modname:
987757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            return self.modname.replace('.', '_')
997757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        else:
1007757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            root = os.path.splitdrive(self.name)[1]
1017757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            return root.replace('\\', '_').replace('/', '_').replace('.', '_')
1027757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
1037757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch    def source_file(self):
1047757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        """Return an open file for reading the source of the code unit."""
1057757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        if os.path.exists(self.filename):
1067757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            # A regular text file: open it.
1077757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            return open_source(self.filename)
1087757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
1097757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        # Maybe it's in a zip file?
1107757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        source = self.file_locator.get_zip_data(self.filename)
1117757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        if source is not None:
1127757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            return StringIO(source)
1137757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch
1147757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        # Couldn't find source.
1157757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch        raise CoverageException(
1167757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            "No source for code %r." % self.filename
1177757ec2eadfa2dd8ac2aeed0a4399e9b07ec38cbBen Murdoch            )
118