"""
Custom test loader class to replace unittest's TestLoader class.
"""

__version__ = "$Revision: 11270 $"

from unittest import TestLoader as _TestLoader

from _testsuite import TestSuite

__all__ = ["TestLoader"]


class TestLoader(_TestLoader):

    """Custom test loader class to replace unittest's TestLoader class.

    This class understands the @test decorators.
    
    """

    suiteClass = TestSuite

    def getTestCaseNames(self, cls):
        def is_test_method(name):
            try:
                attr = getattr(cls, name)
                if not callable(attr):
                    return False
                return self.is_test_method(attr)
            except AttributeError:
                return False
        names = set(_TestLoader.getTestCaseNames(self, cls))
        names.update([name for name in dir(cls) if is_test_method(name)])
        return names

    def is_test_method(self, method):
        return hasattr(method, "_unittest") and method._unittest