Survey code repositories for language feature usage, library usage, and more
"""Library for surveying source code for the frequency of various features.

Typical usage:

```python
from codesurvey import CodeSurvey, LocalSource
from codesurvey.sources import LocalSource
from codesurvey.analyzers.python import PythonAstAnalyzer
from codesurvey.analyzers.python.features import has_set

survey = CodeSurvey(
    sources=[LocalSource(['/path/to/some_python_code'])],
    analyzers=[PythonAstAnalyzer(feature_finders=[has_set])],
)
survey.run(max_repos=1)
print(survey.get_repo_features())
print(survey.get_code_features())
print(survey.get_survey_tree())
```

"""

__version__ = '0.1.3'

from .core import CodeSurvey
from .database import RepoFeature, CodeFeature
from .utils import logger
from . import sources
from . import analyzers

__all__ = [
    'CodeSurvey',
    'RepoFeature',
    'CodeFeature',
    'logger',
    'sources',
    'analyzers',
]

codesurvey/__init__.py
when-of-python/codesurvey
from .core import Analyzer, Code, CodeThunk, FileAnalyzer, FileInfo
from .features import (
    Feature, FeatureDict, FeatureFinder,
    feature_finder, partial_feature_finder, union_feature_finder,
)

__all__ = [
    'Analyzer',
    'Code',
    'CodeThunk',
    'FileAnalyzer',
    'FileInfo',
    'Feature',
    'FeatureDict',
    'FeatureFinder',
    'feature_finder',
    'partial_feature_finder',
    'union_feature_finder',
]

codesurvey/analyzers/__init__.py
when-of-python/codesurvey
"""Base classes for Analyzers of code in Repos."""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import partial
from glob import iglob
import os.path
from typing import Callable, Dict, Generic, Iterator, Optional, Sequence, Union

from codesurvey.utils import get_duplicates
from codesurvey.sources import Repo, TestSource
from .features import CodeReprT, FeatureFinder, Feature


@dataclass(frozen=True)
class Code:
    """Results of analyzing a single unit of source-code from a Repo (e.g.
    a file of source-code) for occurrences of target features."""

    analyzer: 'Analyzer'
    """The Analyzer that performed the analysis."""

    repo: Repo
    """The Repo that the Code belongs to."""

    key: str
    """The unique key of the Code within its Repo."""

    features: Dict[str, Feature]
    """A mapping of feature names to Feature survey results."""


@dataclass(frozen=True)
class CodeThunk:
    """An executable task to be run asynchronously to produce a Code
    analysis."""

    analyzer: 'Analyzer'
    """The Analyzer that will perform the analysis."""

    repo: Repo
    """The Repo that the Code belongs to."""

    key: str
    """The unique key of the Code within its Repo."""

    features: Sequence[str]
    """The names of features to be analyzed."""

    thunk: Callable[[], Code]
    """Function to be called to perform the analysis."""


class Analyzer(ABC, Generic[CodeReprT]):
    """Analyzes Repos to produce Code feature analysis results of
    individual units of source-code (such as source-code files).

    The type argument is the representation of source-code units that
    will be provided to FeatureFinders used with Analyzer instances.

    """

    default_name: str
    """Name to be assigned to Analyzers of this type if a custom name is not
    specified."""

    def __init__(self, *, feature_finders: Sequence[FeatureFinder[CodeReprT]],
                 name: Optional[str] = None):
        """
        Args:
            feature_finders:
                The [FeatureFinders][codesurvey.analyzers.FeatureFinder]
                for analyzing each source-code unit.
            name: Name to identify the Analyzer. If `None`, defaults to the
                Analyzer type's default_name
        """
        self.name = self.default_name if name is None else name
        duplicate_feature_names = get_duplicates([feature.name for feature in feature_finders])
        if duplicate_feature_names:
            duplicate_features_str = ', '.join(duplicate_feature_names)
            raise ValueError((f'Cannot instantiate analyzer "{self}" with duplicate '
                              f'feature names: {duplicate_features_str}. '
                              'Please set a unique name for each feature_finder.'))
        self.feature_finders = {feature.name: feature for feature in feature_finders}

    @abstractmethod
    def prepare_code_representation(self, repo: Repo, code_key: str) -> CodeReprT:
        """Returns a representation of a source-code unit that can be passed
        to the [FeatureFinders][codesurvey.analyzers.FeatureFinder] of
        this Analyzer.

        Args:
            repo: Repo containing the source-code to be analyzed.
            code_key: Unique key of the source-code unit to be analyzed
                within the Repo.

        """

    @abstractmethod
    def code_generator(self, repo: Repo, *,
                       get_code_features: Callable[[str], Sequence[str]]) -> Iterator[Union[Code, CodeThunk]]:
        """Generator yielding [Codes][codesurvey.analyzers.Code] analysis
        results of source-code units within the given Repo, or
        [CodeThunks][codesurvey.analyzers.CodeThunk] that can be
        executed to perform such analyses.

        Args:
            repo: Repo containing the source-code to be analyzed.
            get_code_features: A function that will be called by
                `code_generator()` with a Code's key to get the subset of
                [`get_feature_names()`][codesurvey.analyzers.Analyzer.get_feature_names]
                that should be analyzed for that Code.

        """

    def __str__(self):
        return self.name

    def __repr__(self):
        return f'{self.__class__.__name__}({self})'

    def analyze_code(self, repo: Repo, code_key: str, features: Sequence[str]) -> Code:
        """Produces a [Code][codesurvey.analyzers.Code] analysis for a single
        unit of source-code within a Repo.

        Args:
            repo: Repo containing the source-code to be analyzed.
            code_key: Unique key of the source-code unit to be analyzed
                within the Repo.
            features: Names of features to include in the analysis.
                A subset of the names returned by
                [`get_feature_names()`][codesurvey.analyzers.Analyzer.get_feature_names].

        """
        code_repr = self.prepare_code_representation(repo=repo, code_key=code_key)
        if code_repr is None:
            feature_results = {feature_name: Feature(name=feature_name, ignore=True)
                               for feature_name in features}
        else:
            feature_results = {feature_name: self.feature_finders[feature_name](code_repr)
                               for feature_name in features}
        return self.code(
            repo=repo,
            key=code_key,
            features=feature_results,
        )

    def get_feature_names(self) -> Sequence[str]:
        """Returns the names of all features analyzed by this Analyzer instance."""
        return list(self.feature_finders.keys())

    def code(self, **kwargs) -> Code:
        """Internal helper to generate a Code for this Analyzer. Takes the same
        arguments as Code except for analyzer."""
        return Code(analyzer=self, **kwargs)

    def code_thunk(self, **kwargs) -> CodeThunk:
        """Internal helper to generate a CodeThunk for this Analyzer. Takes
        the same arguments as CodeThunk except for analyzer."""
        return CodeThunk(analyzer=self, **kwargs)


@dataclass(frozen=True)
class FileInfo:
    """Details identifying a source-code file within a Repo."""

    repo: Repo
    """Repo that the file belongs to."""

    rel_path: str
    """Relative path to the file from the Repo directory."""

    @property
    def abs_path(self) -> str:
        """Absolute path to the file."""
        return os.path.join(self.repo.path, self.rel_path)


class FileAnalyzer(Analyzer[CodeReprT]):
    """Base class for Analyzers that analyze each source-code file as the
    target unit of code within a Repo.

    The type argument is the representation of source-code units that
    will be provided to FeatureFinders used with Analyzer instances.

    """

    default_file_glob: str
    """Default glob pattern for finding source-code files. To be assigned
    to FileAnalyzers of this type if a custom glob is not specified."""

    default_file_filters: Sequence[Callable[[FileInfo], bool]] = []
    """Default filters to identify files to exclude from analysis. To be
    assigned to FileAnalyzers of this type if custom filters are not
    specified."""

    def __init__(self, feature_finders: Sequence[FeatureFinder], *,
                 file_glob: Optional[str] = None,
                 file_filters: Optional[Sequence[Callable[[FileInfo], bool]]] = None,
                 name: Optional[str] = None):
        """
        Args:
            feature_finders:
                The [FeatureFinders][codesurvey.analyzers.FeatureFinder]
                for analyzing each source-code file.
            file_glob: Glob pattern for finding source-code files within
                the Repo.
            file_filters: Filters to identify files to exclude from analysis.
                Each filter is a function that takes a
                [`FileInfo`][codesurvey.analyzers.FileInfo] and
                returns `True` if the file should be excluded. file_filters
                cannot be lambdas, as they need to be pickled when passed to
                sub-processes.
            name: Name to identify the Analyzer. If `None`, defaults to the
                Analyzer type's default_name.

        """
        super().__init__(feature_finders=feature_finders, name=name)
        self.file_glob = self.default_file_glob if file_glob is None else file_glob
        self.file_filters = self.default_file_filters if file_filters is None else file_filters

    @abstractmethod
    def prepare_file(self, file_info: FileInfo) -> CodeReprT:
        """Given a [`FileInfo`][codesurvey.analyzers.FileInfo] identifying the
        location of a target source-code file, returns a
        representation of the code that can be passed to the
        [FeatureFinders][codesurvey.analyzers.FeatureFinder] of this
        Analyzer."""

    def prepare_code_representation(self, repo: Repo, code_key: str) -> CodeReprT:
        file_info = FileInfo(repo=repo, rel_path=code_key)
        return self.prepare_file(file_info)

    def _get_file_keys(self, repo: Repo) -> Iterator[str]:
        """Generator yielding the code_keys (relative file paths) within the
        given Repo, applying configured file_filters."""
        for abs_path in iglob(os.path.join(repo.path, self.file_glob), recursive=True):
            if not os.path.isfile(abs_path):
                continue
            file_info = FileInfo(
                repo=repo,
                rel_path=os.path.relpath(abs_path, start=repo.path),
            )
            filtered_out = any([
                file_filter(file_info)
                for file_filter in self.file_filters
            ])
            if filtered_out:
                continue
            yield file_info.rel_path

    def code_generator(self, repo: Repo, *,
                       get_code_features: Callable[[str], Sequence[str]]) -> Iterator[CodeThunk]:
        for file_key in self._get_file_keys(repo):
            features = get_code_features(file_key)
            if len(features) == 0:
                continue

            yield self.code_thunk(
                repo=repo,
                key=file_key,
                features=features,
                thunk=partial(self.analyze_code, repo=repo, code_key=file_key, features=features),
            )

    def test(self, code_snippet: str, *, test_filename: str = 'test_file.txt') -> Dict[str, Feature]:
        """Utility for directly analyzing a string of source-code.

        A Repo will be created in a temporary directory to perform
        analysis of a file created with the given `code_snippet`.

        Args:
            code_snippet: String of source-code to analyze.
            test_filename: Optional custom filename used for the test file.

        Returns:
            A dictionary mapping feature names to
                [`Feature`][codesurvey.analyzers.Feature] results.

        """
        source = TestSource({test_filename: code_snippet})
        repo = next(source.repo_generator())
        code = self.analyze_code(
            repo=repo,
            code_key=test_filename,
            features=self.get_feature_names(),
        )
        return code.features

codesurvey/analyzers/core.py
when-of-python/codesurvey
from dataclasses import dataclass, field
from functools import wraps
from typing import cast, Any, Callable, Dict, Generic, Mapping, Protocol, Sequence, TypeVar, Union


@dataclass
class Feature:
    """Analysis result of a single feature for a single source-code unit."""

    name: str
    """Name of the feature analyzed."""

    occurrences: Sequence[Dict[str, Any]] = field(default_factory=list)
    """Occurrences of the feature within the source-code unit.

    The data captured for each occurrence is determined by the
    [`FeatureFinder`][codesurvey.analyzers.FeatureFinder] that
    produced the Feature.

    """

    ignore: bool = False
    """If `True`, analysis of the feature was skipped for the source-code unit."""


FeatureDict = Dict[str, Any]


CodeReprT = TypeVar('CodeReprT')
# Need a separate contravariant type variable for the protocol below.
CodeReprInputT = TypeVar('CodeReprInputT', contravariant=True)


class FeatureFinder(Protocol, Generic[CodeReprInputT]):
    """Callable for producing a [`Feature`][codesurvey.analyzers.Feature]
    analysis for a given source-code unit representation."""

    name: str
    """Name of the feature that is analyzed."""

    def __call__(self, code_repr: CodeReprInputT) -> Feature:
        """Analyze the given source-code unit representation."""
        pass


def _normalize_feature(*, name: str, feature: Union[Feature, FeatureDict]) -> Feature:
    """Helper for producing a named Feature result from a given Feature or
    FeatureDict.

    Used to allow user-defined feature finders to return either a
    Feature object or equivalent dictionary.

    Args:
        name: Name of the feature
        feature: Feature or dictionary with equivalent data.

    """
    if isinstance(feature, Feature):
        feature.name = name
        return feature
    elif isinstance(feature, Mapping):
        return Feature(name=name, **feature)
    else:
        raise ValueError(('Feature finders must either return a Feature object '
                          'or a dictionary with which to create a Feature. '
                          f'Received: {feature}'))


FeatureFinderFunction = Callable[[CodeReprT], Union[Feature, FeatureDict]]


def feature_finder(name: str) -> Callable[[FeatureFinderFunction[CodeReprT]], FeatureFinder[CodeReprT]]:
    """Decorator for defining a named `FeatureFinder`.

    Example usage:

    ```python
    @feature_finder('while')
    def has_while(code_representation):
        if code_representation is None:
            return {'ignore': True}
        return {
            'occurrences': [
                {'character_index': match.start()}
                for match in re.finditer('test', str(code_representation))
            ]
        }
    ```

    """

    def decorator(func: FeatureFinderFunction[CodeReprT]) -> FeatureFinder[CodeReprT]:

        @wraps(func)
        def decorated(*args, **kwargs):
            feature = func(*args, **kwargs)
            return _normalize_feature(name=name, feature=feature)

        wrapped_feature_finder = cast(FeatureFinder[CodeReprT], decorated)
        wrapped_feature_finder.name = name
        return wrapped_feature_finder

    return decorator


class PartialFeatureFinder(Generic[CodeReprT]):

    def __init__(self, name: str, feature_finder_function: FeatureFinderFunction[CodeReprT], args: Sequence, kwargs: Mapping):
        self.name = name
        self.feature_finder_function = feature_finder_function
        self.args = args
        self.kwargs = kwargs

    def __call__(self, *args, **kwargs):
        keywords = {**self.kwargs, **kwargs}
        feature = self.feature_finder_function(*self.args, *args, **keywords)
        return _normalize_feature(name=self.name, feature=feature)


# Create a constructor function for PartialFeatureFinder for cleaner
# interface documentation. A class must be used to produce pickle-able
# partial objects.
def partial_feature_finder(name: str,
                           feature_finder_function: Callable[..., Union[Feature, FeatureDict]],
                           *args, **kwargs) -> FeatureFinder[CodeReprT]:
    """Defines a FeatureFinder from the partial application of the given
    `feature_finder_function` with the given args and kwargs.

    Example usage:

    ```python
    has_math_module = partial_feature_finder('math_module', module_feature_finder, module='math')
    ```

    """
    return PartialFeatureFinder(
        name=name,
        feature_finder_function=feature_finder_function,
        args=args,
        kwargs=kwargs,
    )


def _union_feature_finder(*args, feature_finders: Sequence[FeatureFinder], **kwargs) -> FeatureDict:
    feature_results = [finder(*args, **kwargs) for finder in feature_finders]
    return dict(
        ignore=all([result.ignore for result in feature_results]),
        occurrences=[
            occurrence
            for result in feature_results
            for occurrence in result.occurrences
        ]
    )


def union_feature_finder(name: str, feature_finders: Sequence[FeatureFinder[CodeReprT]]) -> FeatureFinder[CodeReprT]:
    """Defines a FeatureFinder that returns the union of the occurrences
    of the given `feature_finders`.

    The feature will only return an `'ignore'` result if all
    `feature_finders` ignore the given source-code unit.

    Example usage:

    ```python
    has_loop = union_feature_finder('loop', [has_for_loop, has_while_loop])
    ```

    """
    return partial_feature_finder(name, _union_feature_finder, feature_finders=feature_finders)

codesurvey/analyzers/features.py
when-of-python/codesurvey
from .core import PythonAstAnalyzer
from .features import py_ast_feature_finder, py_module_feature_finder, py_ast_feature_finder_with_transform

__all__ = [
    'PythonAstAnalyzer',
    'py_ast_feature_finder',
    'py_ast_feature_finder_with_transform',
    'py_module_feature_finder',
]

codesurvey/analyzers/python/__init__.py
when-of-python/codesurvey
import ast
from multiprocessing import Process, Queue
import re
from typing import Optional

import astpath
from lxml.etree import Element

from codesurvey.analyzers import FileAnalyzer, FileInfo
from codesurvey.utils import logger

SITE_PACKAGES_REGEX = re.compile(r'.*[/\\]site-packages[/\\].*')


def py_site_packages_filter(file_info):
    """Filter to exclude files under a `site-packages` directory."""
    return SITE_PACKAGES_REGEX.fullmatch(file_info.rel_path)


class PythonAstAnalyzer(FileAnalyzer[Element]):
    """Analyzer that finds .py files and parses them into lxml documents
    representing Python abstract syntax trees for feature analysis."""
    default_name = 'python'
    default_file_glob = '**/*.py'
    default_file_filters = [
        py_site_packages_filter,
    ]
    """Excludes files under a `site-packages` directory that are unlikely
    to belong to the Repo under analysis."""

    def _subprocess_parse_ast(self, *, queue: Queue, file_text: str):
        try:
            file_tree = ast.parse(file_text)
        except Exception as ex:
            queue.put(ex)
        else:
            queue.put(file_tree)

    def prepare_file(self, file_info: FileInfo) -> Optional[Element]:
        with open(file_info.abs_path, 'r') as f:
            file_text = f.read()

        # Parse ast in a subprocess, as sufficiently complex files can
        # crash the interpreter:
        # https://docs.python.org/3/library/ast.html#ast.parse
        queue: Queue = Queue()
        process = Process(target=self._subprocess_parse_ast, kwargs=dict(
            queue=queue,
            file_text=file_text,
        ))
        try:
            process.start()
            result = queue.get()
            process.join()
            if isinstance(result, Exception):
                raise result
        except (SyntaxError, ValueError) as ex:
            logger.error((f'Skipping Python file "{file_info.rel_path}" in '
                          f'repo "{file_info.repo}" that could not be parsed: {ex}'))
            return None

        file_tree = result
        file_xml = astpath.convert_to_xml(file_tree)
        return file_xml

codesurvey/analyzers/python/core.py
when-of-python/codesurvey
from functools import wraps
from typing import Callable, Optional, Sequence

import lxml.etree

from codesurvey.analyzers import (
    FeatureDict, FeatureFinder,
    feature_finder, partial_feature_finder, union_feature_finder,
)


ElementTransform = Callable[[lxml.etree.Element], Optional[lxml.etree.Element]]


def get_first_line_number(element: lxml.etree.Element) -> Optional[int]:
    """Returns the earliest line number of an element of any of its descendents."""
    line_no_strs = set(element.xpath('descendant-or-self::*[@lineno]/@lineno'))
    if len(line_no_strs) == 0:
        return None
    return min(map(int, line_no_strs))


def _py_ast_feature_finder(xml: lxml.etree.Element, *, xpath: str,
                           transform: Optional[ElementTransform] = None) -> FeatureDict:
    elements = xml.xpath(f'descendant-or-self::{xpath}')
    if transform is not None:
        elements = [el for el in map(transform, elements) if el is not None]
    return dict(
        occurrences=[
            {'first_line_number': get_first_line_number(el)}
            for el in elements
        ],
    )


def py_ast_feature_finder(name: str, *, xpath: str) -> FeatureFinder[lxml.etree.Element]:
    """Defines a FeatureFinder that looks for elements in a Python AST
    matching the given xpath query.

    To explore the AST structure of the code constructs you are
    interested in identifying, consider using a tool like:
    https://python-ast-explorer.com/

    """
    return partial_feature_finder(name, _py_ast_feature_finder, xpath=xpath)


def py_ast_feature_finder_with_transform(name: str, *, xpath: str) -> Callable[[ElementTransform], FeatureFinder[lxml.etree.Element]]:
    """Decorator for defining a FeatureFinder that looks for elements in a
    Python AST matching the given xpath query, transforming found elements
    with decorated function.

    The function should receive and return an `lxml.etree.Element`, or
    return `None` if the element should not be considered an
    occurrence of the feature.

    Example usage to look for function calls where the function name
    is 'set':

    ```python
    @py_ast_feature_finder_with_transform('set_function', xpath='Call/func/Name')
    def has_set_function(func_name_el):
        if func_name_el.get('id') == 'set':
            return func_name_el
        return None
    ```

    To explore the AST structure of the code constructs you are
    interested in identifying, consider using a tool like:
    https://python-ast-explorer.com/

    """

    def decorator(func: ElementTransform) -> FeatureFinder[lxml.etree.Element]:

        @wraps(func)
        @feature_finder(name)
        def decorated(xml):
            return _py_ast_feature_finder(xml, xpath=xpath, transform=func)

        return decorated

    return decorator


def _py_module_feature_finder(xml: lxml.etree.Element, *, modules: Sequence[str]):
    matched_els = []

    # import syntax
    import_els = xml.xpath('descendant-or-self::Import')
    for import_el in import_els:
        indiv_import_els = import_el.xpath('names/alias')
        for indiv_import_el in indiv_import_els:
            module = indiv_import_el.get('name')
            if module and any([module.startswith(target_module) for target_module in modules]):
                matched_els.append(indiv_import_el)
    # from syntax
    from_els = xml.xpath('descendant-or-self::ImportFrom')
    for from_el in from_els:
        module = from_el.get('module')
        if module and any([module.startswith(target_module) for target_module in modules]):
            matched_els.append(from_el)

    return dict(
        occurrences=[
            {'first_line_number': get_first_line_number(el)}
            for el in matched_els
        ],
    )


def py_module_feature_finder(name: str, *, modules: Sequence[str]) -> FeatureFinder:
    """Defines a FeatureFinder that looks for import statements of one or
    more target Python `modules`.

    Example usage:

    ```python
    has_dataclasses = py_module_feature_finder('dataclasses_module', modules=['dataclasses'])
    ```

    """
    return partial_feature_finder(name, _py_module_feature_finder, modules=modules)


# ==== Python AST Feature Finders ====

# https://docs.python.org/3/library/ast.html#ast.Try
@py_ast_feature_finder_with_transform('for_else', xpath='For/orelse')
def has_for_else(orelse_el):
    """FeatureFinder for else clauses in for loops."""
    if len(orelse_el) == 0:
        return None
    return orelse_el


# https://docs.python.org/3/library/ast.html#ast.Try
@py_ast_feature_finder_with_transform('try_finally', xpath='Try/finalbody')
def has_try_finally(finalbody_el):
    """FeatureFinder for finally clauses in try statements."""
    if len(finalbody_el) == 0:
        return None
    return finalbody_el


# E.g. `def greeting(name: str):`
@py_ast_feature_finder_with_transform('type_hint', xpath='FunctionDef/args/arguments//annotation')
def has_type_hint(annotation_el):
    """FeatureFinder for type hints."""
    if annotation_el.get('*') is not None:
        return annotation_el
    return None


# The set function
@py_ast_feature_finder_with_transform('set_function', xpath='Call/func/Name')
def has_set_function(func_name_el):
    """FeatureFinder for the set function."""
    if func_name_el.get('id') == 'set':
        return func_name_el
    return None


# A set literal
has_set_value = py_ast_feature_finder('set_value', xpath='Set')
"""FeatureFinder for set literals."""

# The set function or literal
has_set = union_feature_finder('set', [has_set_function, has_set_value])
"""FeatureFinder for sets."""

# Node representing a single formatting field in an f-string.
# https://docs.python.org/3/library/ast.html#ast.FormattedValue
has_fstring = py_ast_feature_finder('fstring', xpath='FormattedValue')
"""FeatureFinder for f-strings."""

# E.g. `if b else c`
# https://docs.python.org/3/library/ast.html#ast.IfExp
has_ternary = py_ast_feature_finder('ternary', xpath='IfExp')
"""FeatureFinder for ternary expressions."""

# https://docs.python.org/3/library/ast.html#ast.Match
has_pattern_matching = py_ast_feature_finder('pattern_matching', xpath='Match')
"""FeatureFinder for pattern matching."""

# https://docs.python.org/3/library/ast.html#ast.NamedExpr
has_walrus = py_ast_feature_finder('walrus', xpath='NamedExpr')
"""FeatureFinder for the walrus operator."""

codesurvey/analyzers/python/features.py
when-of-python/codesurvey
"""Top-level components for running and analyzing code surveys."""

import concurrent.futures
from dataclasses import dataclass
from itertools import cycle
import os
import signal
import sys
from typing import cast, Callable, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple, Union

from tqdm import tqdm
from tqdm.contrib.logging import logging_redirect_tqdm
from tqdm.utils import _screen_shape_wrapper

from .utils import logger, BreakException, get_duplicates, recursive_update
from .sources import Repo, RepoThunk, Source
from .analyzers import Code, CodeThunk, Analyzer
from .database import Database, RepoFeature, CodeFeature


@dataclass(frozen=True)
class Job:
    callback: Callable[[concurrent.futures.Future], None]


@dataclass(frozen=True)
class RepoJob(Job):
    source: Source
    analyzer_features: Dict[str, Sequence[str]]


@dataclass(frozen=True)
class CodeJob(Job):
    analyzer: Analyzer
    repo: Repo


class CodeSurveyRunner:
    """Manages the execution of surveys."""

    def __init__(self, survey: 'CodeSurvey'):
        """
        Args:
            survey: The survey that this runner will execute.
        """
        # Survey attributes
        self.survey = survey
        self.sources = survey.sources
        self.analyzers = survey.analyzers
        self.analyzer_features = survey.analyzer_features
        self.max_workers = survey.max_workers
        self.continue_on_failure = survey.continue_on_failure
        self.save_code_features = survey.save_code_features
        self.save_occurrences = survey.save_occurrences
        self.use_saved_features = survey.use_saved_features

    def get_pbars(self, *, disable_progress: bool,
                  progress_analyzer_features: Optional[Mapping[str, Sequence[str]]],
                  ) -> Dict[Union[str, Tuple[str, ...]], tqdm]:
        """Returns progress bars for the processing of Repos and Codes, as
        well as a border above both progress bars.

        Args:
            disable_progress: If True, progress bar objects will be disabled.
            progress_analyzer_features: See [[codesurvey.core.CodeSurvey.run]]

        """
        if progress_analyzer_features is None:
            all_features = [
                feature
                for features in self.analyzer_features.values()
                for feature in features
            ]
            max_auto_feature_progress = 10
            if len(all_features) > max_auto_feature_progress:
                logger.warning((
                    'Disabling progress trackers for features as there are more '
                    f'than {max_auto_feature_progress} features. To prevent this '
                    'warning, set progress_analyzer_features={}'
                ))
                progress_analyzer_features = {}
            else:
                progress_analyzer_features = self.analyzer_features

        cols = 10
        screen_shape = _screen_shape_wrapper()
        if screen_shape:
            cols = screen_shape(sys.stderr)[0] or cols

        bar_format_with_total = '{desc}: {n_fmt}/{total_fmt} [{elapsed}, {rate_fmt}{postfix}]|{bar}| {percentage:3.0f}% [{remaining} remaining]'
        bar_format_without_total = '{desc}: {n_fmt} [{elapsed}, {rate_fmt}{postfix}]'

        border_pbar = tqdm(
            bar_format=('=' * cols),
            disable=disable_progress,
        )
        repo_pbar = tqdm(
            desc='Repos',
            unit='repos',
            total=self.max_repos,
            bar_format=(bar_format_without_total if self.max_repos is None else bar_format_with_total),
            disable=disable_progress,
        )
        code_pbar = tqdm(
            desc='Codes',
            unit='codes',
            total=self.max_codes,
            bar_format=(bar_format_without_total if self.max_codes is None else bar_format_with_total),
            disable=disable_progress,
        )
        feature_pbars: Dict[Union[str, Tuple[str, ...]], tqdm] = {
            ('feature', analyzer_name, feature_name): tqdm(
                desc=f'Repos with {analyzer_name}:{feature_name}',
                unit='repos',
                bar_format='{desc}: {n_fmt}',
                disable=disable_progress,
            )
            for analyzer_name, feature_names in progress_analyzer_features.items()
            for feature_name in feature_names
        }
        return {
            'border': border_pbar,
            'repos': repo_pbar,
            'codes': code_pbar,
            **feature_pbars,
        }

    def handle_failure(self, *, ex: Exception, message: str):
        """Callback to handle exceptions raised in runner subprocesses."""
        # Simplify subprocess error tracebacks by reporting the cause directly.
        if isinstance(ex.__cause__, concurrent.futures.process._RemoteTraceback):
            ex = ex.__cause__

        if self.continue_on_failure:
            logger.error(f'{message}, skipping: {ex}')
        else:
            logger.error(f'{message}')
            # Simplify traceback by clearing the exception chain.
            raise ex from None

    def get_repo_generator(self) -> Iterator[Union[Repo, RepoThunk]]:
        """Returns a generator that draws Repos/RepoThunks from Sources in a
        round-robin fashion."""
        source_repo_generators = {source: source.repo_generator()
                                  for source in self.sources.values()}
        # Round-robin over Sources.
        for source in cycle(self.sources.values()):
            # Finish when all Sources are exhausted
            if len(source_repo_generators) == 0:
                return
            # Skip to the next Source if this one is exhausted
            if source not in source_repo_generators:
                continue
            try:
                yield next(source_repo_generators[source])
            except StopIteration:
                # Remove exhausted Sources
                del source_repo_generators[source]
            except Exception as ex:
                self.handle_failure(
                    ex=ex,
                    message=f'Failed to fetch repo from source "{source}"',
                )

    def get_executor(self) -> concurrent.futures.ProcessPoolExecutor:
        """Returns a multi-processing executor for running repo-fetching and
        code-processing subprocesses."""

        def init_worker():
            # Ignore keyboard interrupts in subprocesses.
            signal.signal(signal.SIGINT, signal.SIG_IGN)

        return concurrent.futures.ProcessPoolExecutor(
            max_workers=self.max_workers,
            initializer=init_worker,
        )

    def reached_max_repos(self) -> bool:
        """Checks whether the number of Repos being fetched, being analyzed,
        and fully analyzed has reached the configured max_repos."""
        if self.max_repos is None:
            return False

        pending_repos = len([job for job in self.future_to_job.values()
                             if isinstance(job, RepoJob)])
        repo_count = len(self.current_repos) + pending_repos + self.completed_repo_count
        return repo_count >= self.max_repos

    def reached_max_codes(self) -> bool:
        """Checks whether the number of analyzed Codes has reached the
        configured max_repos."""
        if self.max_codes is None:
            return False

        return self.completed_code_count >= self.max_codes

    def check_repo_completion(self, repos: Sequence[Repo]) -> None:
        """Check whether there are any completed current_repos without any
        remaining CodeJobs, save their results, and remove them from
        current_repos."""
        completed_repos = [
            repo for repo in self.current_repos
            if len([
                job for job in self.future_to_job.values()
                if isinstance(job, CodeJob) and job.repo is repo
            ]) == 0
        ]
        for repo in completed_repos:
            self.db.save_repo_features(repo, keep_code_features=self.save_code_features)
            self.completed_repo_count += 1
            self.pbars['repos'].update(1)

            repo_features = self.db.get_repo_features(
                source_names=[repo.source.name],
                repo_keys=[repo.key],
            )
            for repo_feature in repo_features:
                if repo_feature.occurrence_count > 0:
                    feature_pbar_key = (
                        'feature',
                        repo_feature.analyzer_name,
                        repo_feature.feature_name,
                    )
                    self.pbars[feature_pbar_key].update(1)

            logger.info(f'Completed repo "{repo}"')
            repo.cleanup()
            self.current_repos.remove(repo)

    def handle_code(self, *, code: Code) -> None:
        """Save survey results for the given Code, updating progress tracking."""
        if self.reached_max_codes():
            return

        self.db.save_code_features(code, save_occurrences=self.save_occurrences)
        self.pbars['codes'].update(1)
        self.completed_code_count += 1

    def handle_code_future(self, future: concurrent.futures.Future) -> None:
        """Handle saving a completed CodeJob, handling Job failure."""
        job = cast(CodeJob, self.future_to_job[future])
        try:
            code = future.result()
        except Exception as ex:
            self.handle_failure(
                ex=ex,
                message=f'Failed to analyze code from repo "{job.repo}" with analyzer "{job.analyzer}"',
            )
        else:
            self.handle_code(code=code)

    def handle_repo(self, *, repo: Repo, analyzer_features: Mapping[str, Sequence[str]]) -> None:
        """Analyze the given Repo for the given analyzer_features.

        All Codes for the Repo will either be analyzed or have a
        pending CodeJob after calling this function.

        """
        # Add the repo to the list of current_repos currently being analyzed.
        self.current_repos.append(repo)
        self.db.save_repo_metadata(repo)
        try:
            if self.reached_max_codes():
                raise BreakException()
            for analyzer in self.analyzers.values():
                features = analyzer_features.get(analyzer.name)
                if features is None:
                    logger.info(f'Skipping completed analyzer "{analyzer}" for repo "{repo}"')
                    continue

                analyzer = self.analyzers[analyzer.name]
                try:
                    logger.info(f'Analyzing repo "{repo}" with analyzer "{analyzer}"')

                    def get_code_features(code_key: str) -> Sequence[str]:
                        """Determine which features still need to be surveyed for the given
                        Code (all features if not use_saved_features)"""
                        if not self.use_saved_features:
                            return cast(Sequence[str], features)
                        return self.db.get_code_missing_features(
                            source_name=repo.source.name,
                            repo_key=repo.key,
                            analyzer_name=analyzer.name,
                            code_key=code_key,
                            features=features,
                        )

                    analyzer_codes = analyzer.code_generator(
                        repo=repo,
                        get_code_features=get_code_features,
                    )
                except Exception as ex:
                    self.handle_failure(
                        ex=ex,
                        message=f'Failed to get analyzer "{analyzer}" codes for repo "{repo}"',
                    )
                else:
                    # Loop over each Code found by the Analyzer
                    for code_or_thunk in analyzer_codes:
                        if self.reached_max_codes():
                            raise BreakException()

                        if isinstance(code_or_thunk, CodeThunk):
                            # If it's a CodeThunk, submit a Job to
                            # analyze the Code, with a callback to
                            # save the results.
                            future = self.executor.submit(cast(Callable, code_or_thunk.thunk))
                            self.future_to_job[future] = CodeJob(
                                analyzer=analyzer,
                                repo=repo,
                                callback=self.handle_code_future,
                            )
                        else:
                            # If it's a Code that is already analyzed,
                            # go straight to saving the results.
                            self.handle_code(code=cast(Code, code_or_thunk))
        except BreakException:
            logger.info(f'Max codes reached, "{repo}" will not be fully analyzed')
        finally:
            self.check_repo_completion([repo])

    def handle_repo_future(self, future: concurrent.futures.Future) -> None:
        """Handle analysis of a completed RepoJob, handling Job failure."""
        job = cast(RepoJob, self.future_to_job[future])
        try:
            repo = future.result()
        except Exception as ex:
            self.handle_failure(
                ex=ex,
                message=f'Failed to fetch repo from source "{job.source}"',
            )
        else:
            self.handle_repo(
                repo=repo,
                analyzer_features=job.analyzer_features,
            )

    def consume_repos(self) -> None:
        """Consume Repos from the repo_generator

        When a Source yields a RepoThunk or an Analyzer yields a
        CodeThunk, it starts a RepoJob/CodeJob and adds it to
        future_to_job - and this function returns when the number of
        pending tasks reaches the number of sub-process workers.

        When a Source yields a Repo or an Analyzer yields a
        Code, they are handled synchronously in this function.

        """
        while (
            # Check for idle workers
            (len(self.future_to_job) < self.max_workers)
            # Check for max_repos and max_codes
            and not self.reached_max_repos()
            and not self.reached_max_codes()
        ):
            # Get the next Repo or RepoThunk from the repo_generator
            try:
                repo_or_thunk = next(self.repo_generator)
            except StopIteration:
                break
            # Determine which Analyzer features still need to be
            # surveyed for the Repo (all features if not
            # use_saved_features)
            repo_analyzer_features = (
                self.db.get_repo_missing_analyzer_features(
                    source_name=repo_or_thunk.source.name,
                    repo_key=repo_or_thunk.key,
                    analyzer_features=self.analyzer_features,
                )
                if self.use_saved_features
                else self.analyzer_features
            )

            # Continue if this is a duplicate of a repo we're already analyzing.
            if any((repo.source.name == repo_or_thunk.source.name)
                   and (repo.key == repo_or_thunk.key)
                   for repo in self.current_repos):
                logger.info(f'Skipping in-progress repo "{repo_or_thunk}"')
                continue

            # Continue if there are no missing features.
            if len(repo_analyzer_features) == 0:
                logger.info(f'Skipping fully analyzed repo "{repo_or_thunk}"')
                continue

            if isinstance(repo_or_thunk, RepoThunk):
                # If it's a RepoThunk, submit a Job to fetch the Repo,
                # with a callback to handle it's analysis later.
                logger.info(f'Fetching repo "{repo_or_thunk}"')
                future = self.executor.submit(cast(Callable, repo_or_thunk.thunk))
                self.future_to_job[future] = RepoJob(
                    source=repo_or_thunk.source,
                    analyzer_features=repo_analyzer_features,
                    callback=self.handle_repo_future,
                )
            else:
                # If it's a Repo that doesn't need to be fetched, go
                # straight to handling the analysis of the Repo.
                self.handle_repo(
                    repo=cast(Repo, repo_or_thunk),
                    analyzer_features=repo_analyzer_features,
                )

    def run(self,
            max_repos: Optional[int] = None,
            max_codes: Optional[int] = None,
            disable_progress: bool = False,
            progress_analyzer_features: Optional[Mapping[str, Sequence[str]]] = None) -> None:
        """Start the runner."""
        # Run configuration
        self.max_repos = max_repos
        self.max_codes = max_codes
        # State initialization
        self.completed_repo_count = 0
        self.completed_code_count = 0
        self.db = self.survey.get_db()
        self.pbars = self.get_pbars(
            disable_progress=disable_progress,
            progress_analyzer_features=progress_analyzer_features,
        )
        self.repo_generator = self.get_repo_generator()
        self.executor = self.get_executor()
        # Keep track of jobs submitted to the executor
        self.future_to_job: Dict[concurrent.futures.Future, Job] = {}
        # Keep track of Repos that are currently being analyzed.
        self.current_repos: List[Repo] = []

        logger.info(f'Preparing database in {self.db.filepath}')
        self.db.initialize()

        with logging_redirect_tqdm(loggers=[logger]):
            try:
                while True:
                    # Start Repo-fetching and Code-analyzing jobs and
                    # add them to future_to_job until all workers are
                    # occupied.
                    self.consume_repos()

                    # If there are no jobs left uncompleted or unprocessed,
                    if len(self.future_to_job) == 0:
                        break

                    # Wait for at least one job to finish.
                    done, _ = concurrent.futures.wait(
                        self.future_to_job,
                        return_when=concurrent.futures.FIRST_COMPLETED,
                    )
                    # Process each finished job
                    for future in done:
                        job = self.future_to_job[future]
                        job.callback(future)
                    # Remove finished sub-tasks from future_to_job
                    self.future_to_job = {future: self.future_to_job[future]
                                          for future in self.future_to_job
                                          if future not in done}
                    # Finish the processing of any completed
                    # current_repos, removing them from current_repos.
                    self.check_repo_completion(self.current_repos)
            except KeyboardInterrupt:
                logger.info('Interrupted')
                raise
            finally:
                for future in self.future_to_job:
                    future.cancel()
                for process in self.executor._processes.values():
                    process.terminate()
                self.executor.shutdown()
                for repo in self.current_repos:
                    repo.cleanup()
                for pbar in self.pbars.values():
                    pbar.close()
                self.db.close()


class CodeSurvey:
    """Primary interface for running surveys and inspecting their results.

    A CodeSurvey is instantiated with a set of
    [Sources](sources/core.md) to be surveyed and a set of
    [Analyzers](analyzers/core.md) to count the occurrences of
    *features* within them. Each Source may fetch multiple *Repos*
    (e.g. a project directory, a git repository), each of which may
    contain multiple *Codes* (e.g. a source-code file) to be analyzed.
    Each Analyzer will be configured to identify a particular set of
    *features* within each Code.

    Additional arguments can be passed to
    [`__init__()`][codesurvey.CodeSurvey.__init__] to control
    persistent storage, parallelism, and other options.

    The survey can be executed with
    [`run()`][codesurvey.CodeSurvey.run], which accepts options that
    determine the stopping condition for the survey. Multiple calls to
    [`run()`][codesurvey.CodeSurvey.run] will extend the results of
    the survey.

    [`get_repo_features()`][codesurvey.CodeSurvey.get_repo_features],
    [`get_code_features()`][codesurvey.CodeSurvey.get_code_features],
    and [`get_survey_tree()`][codesurvey.CodeSurvey.get_survey_tree]
    can be used to inspect the results of the survey.

    Previous survey results can be loaded for inspection specifying
    the same `db_filepath` used for previous survey run(s).

    """

    def __init__(self, *,
                 sources: Sequence[Source],
                 analyzers: Sequence[Analyzer],
                 db_filepath: str = ':memory:',
                 max_workers: Optional[int] = 1,
                 continue_on_failure: bool = True,
                 save_code_features: bool = True,
                 save_occurrences: bool = True,
                 use_saved_features: bool = True):
        """
        Args:
            sources: Sources from which to fetch Repos of Codes
                to survey. If multiple Sources are provided, Repo fetching
                will cycle through them in a round-robin fashion.
            analyzers: Analyzers to identify features in fetched code.
            db_filepath: Path to an sqlite database file for persisting survey
                results. Creates a new sqlite database if the path does not
                exist. Defaults to a non-persistent in-memory database.
            max_workers: The maximum number of parallel worker processes for
                fetching Repos from Sources and executing Analyzers. Defaults
                to a single worker.
            continue_on_failure: If `True`, exceptions raised by Sources and
                Analyzers will be logged, but will not halt the survey.
            save_code_features: If `True`, features of individual Codes will be
                retained in the survey database. Otherwise, Code features will
                be deleted once they have been used to compute aggregate
                features of its respective Repo.
            save_occurrences: If `True`, occurrence objects returned by
                FeatureFinders will be saved in the survey database.
            use_saved_features: If `True`, re-use saved features from an
                Analyzer for a Code when they already exist in the survey
                database. Otherwise, reapply all Analyzers to all Codes.

        Raises:
            ValueError: Invalid survey configuration was specified.

        """
        duplicate_source_names = get_duplicates([source.name for source in sources])
        if duplicate_source_names:
            duplicate_sources_str = ', '.join(duplicate_source_names)
            raise ValueError(('Cannot instantiate CodeSurvey with duplicate '
                              f'source names: {duplicate_sources_str}. '
                              'Please set a unique name for each source.'))
        self.sources = {source.name: source for source in sources}

        duplicate_analyzer_names = get_duplicates([analyzer.name for analyzer in analyzers])
        if duplicate_analyzer_names:
            duplicate_analyzers_str = ', '.join(duplicate_analyzer_names)
            raise ValueError(('Cannot instantiate CodeSurvey with duplicate '
                              f'analyzer names: {duplicate_analyzers_str}. '
                              'Please set a unique name for each analyzer.'))
        self.analyzers = {analyzer.name: analyzer for analyzer in analyzers}

        self.analyzer_features = {analyzer.name: analyzer.get_feature_names()
                                  for analyzer in analyzers}
        self.db_filepath = db_filepath
        self.max_workers = max_workers or os.cpu_count() or 1
        self.continue_on_failure = continue_on_failure
        self.save_code_features = save_code_features
        self.save_occurrences = save_occurrences
        self.use_saved_features = use_saved_features

        self._runner = CodeSurveyRunner(self)

    def get_db(self):
        """Returns the Database that persists survey results."""
        return Database(self.db_filepath)

    def run(self, *,
            max_repos: Optional[int] = None,
            max_codes: Optional[int] = None,
            disable_progress: bool = False,
            progress_analyzer_features: Optional[Mapping[str, Sequence[str]]] = None):
        """Runs the survey by fetching code from sources and applying analyzers.

        If neither of the `max_repos` nor `max_codes` stopping
        conditions is specified, the survey will continue running
        until a `KeyboardInterrupt` exception.

        Args:
            max_repos: If specified, the run will stop after analysing this
                many Repos.
            max_codes: If specified, the run will stop after analysing this
                many Codes.
            disable_progress: If `True`, do not display tqdm progress bars
                counting Repos and Codes analyzed.
            progress_analyzer_features: Mapping of analyzer names to sequences
                of feature names for which progress trackers should be
                displayed to count Repos found with those features. Defaults
                to all features, but disables feature progress trackers with
                a warning when there are more than 10 features.

        """
        self._runner.run(
            max_repos=max_repos,
            max_codes=max_codes,
            disable_progress=disable_progress,
            progress_analyzer_features=progress_analyzer_features,
        )

    def get_repo_features(self, *,
                          source_names: Optional[Sequence[str]] = None,
                          analyzer_names: Optional[Sequence[str]] = None,
                          feature_names: Optional[Sequence[str]] = None) -> List[RepoFeature]:
        """Returns RepoFeatures of surveyed Repos.

        Args:
            source_names: If specified, only features from the named Sources
                will be returned.
            analyzer_names: If specified, only features from the named Analyzers
                will be returned.
            feature_names: If specified, only results for the named features
                will be returned.

        """
        return self.get_db().get_repo_features(source_names=source_names,
                                               analyzer_names=analyzer_names,
                                               feature_names=feature_names)

    def get_code_features(self, *,
                          source_names: Optional[Sequence[str]] = None,
                          analyzer_names: Optional[Sequence[str]] = None,
                          feature_names: Optional[Sequence[str]] = None) -> List[CodeFeature]:
        """Returns CodeFeatures of surveyed Codes.

        Only returns results from runs where `save_code_results` was `True`.

        Args:
            source_names: If specified, only features from the named Sources
                will be returned.
            analyzer_names: If specified, only features from the named Analyzers
                will be returned.
            feature_names: If specified, only results for the named features
                will be returned.

        """
        return self.get_db().get_code_features(source_names=source_names,
                                               analyzer_names=analyzer_names,
                                               feature_names=feature_names)

    def get_survey_tree(self, *,
                        source_names: Optional[Sequence[str]] = None,
                        analyzer_names: Optional[Sequence[str]] = None,
                        feature_names: Optional[Sequence[str]] = None) -> Dict:
        """Returns surveyed CodeFeatures and RepoFeatures structured under a
        tree structure of Sources, Repos, and Analyzers.

        Args:
            source_names: If specified, only features from the named Sources
                will be returned.
            analyzer_names: If specified, only features from the named Analyzers
                will be returned.
            feature_names: If specified, only results for the named features
                will be returned.

        Returns:
            A dictionary with the following structure:
                ```python
                {
                    'sources': {
                        '<source_name>': {
                            'repos: {
                                '<repo_key>': {
                                    'analyzers': {
                                        '<analyzer_name>': {
                                            'features': {
                                                'updated': datetime(...),
                                                'occurence_count': int(...),
                                                'code_occurrence_count': int(...),
                                                'code_total_count': int(...),
                                            },
                                            # 'codes' key is only present if
                                            # survey runs are performed with
                                            # `save_code_features=True`
                                            'codes': {
                                                '<code_key>': {
                                                    'features': {
                                                        '<feature_name>': {
                                                            'updated': datetime(...),
                                                            'occurence_count': int(...),
                                                        },
                                                        ...
                                                    }
                                                },
                                                ...
                                            }
                                        },
                                        ...
                                    },
                                    'repo_metadata': {
                                        '<metadata_key>': ...,
                                        ...
                                    }
                                },
                                ...
                            }
                        },
                        ...
                    }
                }
                ```

        """
        tree: Dict = {'sources': {}}
        code_features = self.get_code_features(source_names=source_names,
                                               analyzer_names=analyzer_names,
                                               feature_names=feature_names)
        for c in code_features:
            recursive_update(tree, {
                'sources': {c.source_name: {
                    'repos': {c.repo_key: {
                        'analyzers': {c.analyzer_name: {
                            'codes': {c.code_key: {
                                'features': {c.feature_name: {
                                    'updated': c.updated,
                                    'occurrence_count': c.occurrence_count,
                                    'occurrences': c.occurrences,
                                }}
                            }}
                        }},
                        'repo_metadata': c.repo_metadata,
                    }}
                }}
            })
        repo_features = self.get_repo_features(source_names=source_names,
                                               analyzer_names=analyzer_names)
        for r in repo_features:
            recursive_update(tree, {
                'sources': {r.source_name: {
                    'repos': {r.repo_key: {
                        'analyzers': {r.analyzer_name: {
                            'features': {r.feature_name: {
                                'updated': r.updated,
                                'occurrence_count': r.occurrence_count,
                                'code_occurrence_count': r.code_occurrence_count,
                                'code_total_count': r.code_total_count,
                            }},
                        }},
                        'repo_metadata': r.repo_metadata,
                    }}
                }}
            })
        return tree

codesurvey/core.py
when-of-python/codesurvey
"""Database interaction layer for storing and retrieving survey results."""

from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, List, Mapping, Optional, Protocol, Sequence, Tuple

from playhouse.sqlite_ext import (
    SqliteExtDatabase,
    Model,
    CharField,
    TimestampField,
    IntegerField,
    JSONField,
    CompositeKey,
    fn,
)

from codesurvey.sources import Repo
from codesurvey.analyzers import Code


@dataclass(frozen=True)
class RepoFeature:
    updated: datetime
    """Timestamp when this analysis was last updated."""

    source_name: str
    """Name of the Source that produced the target Repo."""

    repo_key: str
    """Key identifying the target Repo within the Source."""

    analyzer_name: str
    """Name of the Analyzer that produced this feature."""

    feature_name: str
    """Name of the analyzed feature."""

    occurrence_count: int
    """Number of occurrences of this feature within the Repo."""

    code_occurrence_count: int
    """Number of Codes within the Repo containing this feature."""

    code_total_count: int
    """Total number of Codes analyzed for this feature within the Repo."""

    repo_metadata: Dict[str, Any]
    """Metadata of the Repo provided by the Source."""


@dataclass(frozen=True)
class CodeFeature:
    updated: datetime
    """Timestamp when this analysis was last updated."""

    source_name: str
    """Name of the Source that produced the target Repo."""

    repo_key: str
    """Key identifying the target Repo within the Source."""

    analyzer_name: str
    """Name of the Analyzer that produced this feature."""

    code_key: str
    """Key idenfitying the target Code within the Repo."""

    feature_name: str
    """Name of the analyzed feature."""

    occurrence_count: Optional[int]
    """Number of occurrences of this feature within the Code, or `None` if
    analysis of this Code was skipped."""

    occurrences: Optional[List[Dict[str, Any]]]
    """Original occurrence objects returned by FeatureFinders."""

    repo_metadata: Dict[str, Any]
    """Metadata of the Repo provided by the Source."""


class MetadataCache(Protocol):
    """Memoized function for retrieving Repo metadata."""

    def __call__(self, *, source_name: str, repo_key: str) -> Dict[str, Any]:
        pass


class Database:
    """Stores and fetches survey results from an sqlite database."""

    def __init__(self, filepath: str):
        self.filepath = filepath
        self.db = SqliteExtDatabase(self.filepath)

        class BaseModel(Model):
            class Meta:
                database = self.db

            @classmethod
            def get_primary_key_columns(cls):
                return [
                    getattr(cls, column_name)
                    for column_name in self.db.get_primary_keys(cls._meta.table_name)
                ]

        class RepoMetadataModel(BaseModel):
            class Meta:
                table_name = 'repo_metadata'
                primary_key = CompositeKey(
                    'source_name',
                    'repo_key',
                    'metadata_key',
                )

            updated = TimestampField()
            source_name = CharField()
            repo_key = CharField()
            metadata_key = CharField()
            metadata_value = JSONField()

        class RepoFeatureModel(BaseModel):
            class Meta:
                table_name = 'repo_feature'
                primary_key = CompositeKey(
                    'source_name',
                    'repo_key',
                    'analyzer_name',
                    'feature_name',
                )

            updated = TimestampField()
            source_name = CharField()
            repo_key = CharField()
            analyzer_name = CharField()
            feature_name = CharField()
            occurrence_count = IntegerField()
            code_occurrence_count = IntegerField()
            code_total_count = IntegerField()

        class CodeFeatureModel(BaseModel):
            class Meta:
                table_name = 'code_feature'
                primary_key = CompositeKey(
                    'source_name',
                    'repo_key',
                    'analyzer_name',
                    'code_key',
                    'feature_name',
                )

            updated = TimestampField()
            source_name = CharField()
            repo_key = CharField()
            analyzer_name = CharField()
            code_key = CharField()
            feature_name = CharField()
            # null occurrence_count or occurrences indicates an ignored code.
            occurrence_count = IntegerField(null=True)
            occurrences = JSONField(null=True)

        self.RepoMetadataModel = RepoMetadataModel
        self.RepoFeatureModel = RepoFeatureModel
        self.CodeFeatureModel = CodeFeatureModel
        self.tables = [self.RepoMetadataModel, self.RepoFeatureModel, self.CodeFeatureModel]

    def initialize(self):
        """Connect to the database and initialize the schema."""
        self.db.connect()
        self.db.create_tables(self.tables)

    def close(self):
        """Close the database."""
        self.db.close()

    def get_repo_missing_analyzer_features(
        self, *,
        source_name: str,
        repo_key: str,
        analyzer_features: Mapping[str, Sequence[str]],
    ) -> Dict[str, List[str]]:
        """Filter the given analyzer_features to those not recorded for a
        given Repo.

        Args:
            source_name: Name of the target Repo's Source
            repo_key: Key of the target Repo
            analyzer_features: Mapping of Analyzer names to feature names

        """
        rows = (self.RepoFeatureModel
                .select(
                    self.RepoFeatureModel.analyzer_name,
                    self.RepoFeatureModel.feature_name,
                )
                .distinct(True)
                .where(
                    (self.RepoFeatureModel.source_name == source_name)
                    & (self.RepoFeatureModel.repo_key == repo_key)
                    & (self.RepoFeatureModel.analyzer_name.in_(list(analyzer_features.keys())))))
        existing_analyzer_features = defaultdict(set)
        for row in rows:
            existing_analyzer_features[row.analyzer_name].add(row.feature_name)

        missing_analyzer_features = {}
        for analyzer_name, features in analyzer_features.items():
            missing_features = [
                feature for feature in features
                if feature not in existing_analyzer_features[analyzer_name]
            ]
            if len(missing_features) > 0:
                missing_analyzer_features[analyzer_name] = missing_features
        return missing_analyzer_features

    def get_code_missing_features(
        self, *,
        source_name: str,
        repo_key: str,
        code_key: str,
        analyzer_name: str,
        features: Sequence[str],
    ) -> Sequence[str]:
        """Filter the given features to those not recorded for a given Code.

        Args:
            source_name: Name of the target Repo's Source
            repo_key: Key of the target Repo
            code_key: Key of the target Code
            analyzer_name: Name of the Analyzer to get features for
            features: Names of the Analyzer's features

        """
        rows = (self.CodeFeatureModel
                .select(self.CodeFeatureModel.feature_name)
                .distinct(True)
                .where(
                    (self.CodeFeatureModel.source_name == source_name)
                    & (self.CodeFeatureModel.repo_key == repo_key)
                    & (self.CodeFeatureModel.analyzer_name == analyzer_name)
                    & (self.CodeFeatureModel.code_key == code_key)
                    & (self.CodeFeatureModel.feature_name.in_(features))))
        existing_features = set([row.feature_name for row in rows])
        return [
            feature for feature in features
            if feature not in existing_features
        ]

    def save_repo_metadata(self, repo: Repo):
        """Save metadata for the given Repo."""
        with self.db.atomic():
            (self.RepoMetadataModel
             .insert_many([
                 dict(
                     updated=datetime.now(),
                     source_name=repo.source.name,
                     repo_key=repo.key,
                     metadata_key=metadata_key,
                     metadata_value=metadata_value,
                 )
                 for metadata_key, metadata_value in repo.metadata.items()
             ])
             .on_conflict(
                 conflict_target=self.RepoMetadataModel.get_primary_key_columns(),
                 preserve=[
                     self.RepoMetadataModel.updated,
                     self.RepoMetadataModel.metadata_value,
                 ],
             )
             .execute())

    def save_code_features(self, code: Code, *, save_occurrences: bool):
        """Save Analyzer features for the given Code.

        Only save the raw occurrence objects if save_occurrences is
        `True`.

        """
        with self.db.atomic():
            (self.CodeFeatureModel
             .insert_many([
                 dict(
                     updated=datetime.now(),
                     source_name=code.repo.source.name,
                     repo_key=code.repo.key,
                     analyzer_name=code.analyzer.name,
                     code_key=code.key,
                     feature_name=feature_name,
                     occurrence_count=(
                         None if feature.ignore else len(feature.occurrences)
                     ),
                     occurrences=(
                         None if (feature.ignore or not save_occurrences)
                         else feature.occurrences
                     ),
                 )
                 for feature_name, feature in code.features.items()
             ])
             .on_conflict(
                 conflict_target=self.CodeFeatureModel.get_primary_key_columns(),
                 preserve=[
                     self.CodeFeatureModel.updated,
                     self.CodeFeatureModel.occurrence_count,
                     self.CodeFeatureModel.occurrences,
                 ],
             )
             .execute())

    def save_repo_features(self, repo: Repo, *, keep_code_features: bool):
        """Save Analyzer features for the given Repo by aggregating Code
        features.

        If keep_code_features=False, delete the Code features for the
        Repo after saving Analyzer features.

        """
        repo_code_filter = ((self.CodeFeatureModel.source_name == repo.source.name)
                            & (self.CodeFeatureModel.repo_key == repo.key))
        (self.RepoFeatureModel
         .insert_from(
             query=(self.CodeFeatureModel
                    .select(
                        # Most recent CodeFeatureModel updated time
                        fn.MAX(self.CodeFeatureModel.updated),
                        self.CodeFeatureModel.source_name,
                        self.CodeFeatureModel.repo_key,
                        self.CodeFeatureModel.analyzer_name,
                        self.CodeFeatureModel.feature_name,
                        # Sum of all occurrences in CodeFeatureModels
                        fn.SUM(self.CodeFeatureModel.occurrence_count),
                        # Count of all CodeFeatureModels with at least one occurrence
                        fn.SUM(fn.MIN(self.CodeFeatureModel.occurrence_count, 1)),
                        # Count of all CodeFeatureModels
                        fn.COUNT(),
                    )
                    .where(repo_code_filter
                           # Do not count "ignored" CodeFeatureModels
                           & self.CodeFeatureModel.occurrence_count.is_null(False))
                    .group_by(self.CodeFeatureModel.feature_name)),
             fields=[
                 self.RepoFeatureModel.updated,
                 self.RepoFeatureModel.source_name,
                 self.RepoFeatureModel.repo_key,
                 self.RepoFeatureModel.analyzer_name,
                 self.RepoFeatureModel.feature_name,
                 self.RepoFeatureModel.occurrence_count,
                 self.RepoFeatureModel.code_occurrence_count,
                 self.RepoFeatureModel.code_total_count,
             ],
         )
         .on_conflict(
             conflict_target=self.RepoFeatureModel.get_primary_key_columns(),
             preserve=[
                 self.RepoFeatureModel.updated,
                 self.RepoFeatureModel.occurrence_count,
                 self.RepoFeatureModel.code_occurrence_count,
                 self.RepoFeatureModel.code_total_count,
             ],
         )
         .execute())

        # Optionally delete matched CodeFeatureModel rows to reduce storage.
        if not keep_code_features:
            (self.CodeFeatureModel
             .delete()
             .where(repo_code_filter)
             .execute())

    def _get_repo_metadata(self, *, source_name: str, repo_key: str) -> Dict[str, Any]:
        """Retrieves metadata for a given Repo from the database."""
        rows = (self.RepoMetadataModel
                .select(
                    self.RepoMetadataModel.metadata_key,
                    self.RepoMetadataModel.metadata_value,
                )
                .where((self.RepoMetadataModel.source_name == source_name)
                       & (self.RepoMetadataModel.repo_key == repo_key)))
        return {
            row.metadata_key: row.metadata_value
            for row in rows
        }

    def _get_repo_metadata_cache(self) -> MetadataCache:
        """Returns a function that returns the metadata for a given Repo with
        caching to prevent repeated database queries."""
        cache: Dict[Tuple[str, str], Dict[str, Any]] = {}

        def metadata_cache(*, source_name: str, repo_key: str) -> Dict[str, Any]:
            cache_key = (source_name, repo_key)
            if cache_key not in cache:
                cache[cache_key] = self._get_repo_metadata(source_name=source_name, repo_key=repo_key)
            return cache[cache_key]

        return metadata_cache

    def get_repo_features(self, *,
                          source_names: Optional[Sequence[str]] = None,
                          repo_keys: Optional[Sequence[str]] = None,
                          analyzer_names: Optional[Sequence[str]] = None,
                          feature_names: Optional[Sequence[str]] = None) -> List[RepoFeature]:
        """Returns RepoFeatures of surveyed Repos.

        Args:
            source_names: If specified, only features from the named Sources
                will be returned.
            repo_keys: If specified, only features from the named Repos
                will be returned.
            analyzer_names: If specified, only features from the named Analyzers
                will be returned.
            feature_names: If specified, only results for the named features
                will be returned.

        """
        query = self.RepoFeatureModel.select()
        if source_names is not None:
            query = query.where(self.RepoFeatureModel.source_name.in_(source_names))
        if repo_keys is not None:
            query = query.where(self.RepoFeatureModel.repo_key.in_(repo_keys))
        if analyzer_names is not None:
            query = query.where(self.RepoFeatureModel.analyzer_name.in_(analyzer_names))
        if feature_names is not None:
            query = query.where(self.RepoFeatureModel.feature_name.in_(feature_names))

        metadata_cache = self._get_repo_metadata_cache()
        return [
            RepoFeature(
                **row,
                repo_metadata=metadata_cache(source_name=row['source_name'],
                                             repo_key=row['repo_key']),
            )
            for row in query.dicts()
        ]

    def get_code_features(self, *,
                          source_names: Optional[Sequence[str]] = None,
                          repo_keys: Optional[Sequence[str]] = None,
                          analyzer_names: Optional[Sequence[str]] = None,
                          feature_names: Optional[Sequence[str]] = None) -> List[CodeFeature]:
        """Returns CodeFeatures of surveyed Codes.

        Args:
            source_names: If specified, only features from the named Sources
                will be returned.
            repo_keys: If specified, only features from the named Repos
                will be returned.
            analyzer_names: If specified, only features from the named Analyzers
                will be returned.
            feature_names: If specified, only results for the named features
                will be returned.

        """
        query = self.CodeFeatureModel.select()
        if source_names is not None:
            query = query.where(self.CodeFeatureModel.source_name.in_(source_names))
        if repo_keys is not None:
            query = query.where(self.CodeFeatureModel.repo_key.in_(repo_keys))
        if analyzer_names is not None:
            query = query.where(self.CodeFeatureModel.analyzer_name.in_(analyzer_names))
        if feature_names is not None:
            query = query.where(self.CodeFeatureModel.feature_name.in_(feature_names))

        metadata_cache = self._get_repo_metadata_cache()
        return [
            CodeFeature(
                **row,
                repo_metadata=metadata_cache(source_name=row['source_name'],
                                             repo_key=row['repo_key']),
            )
            for row in query.dicts()
        ]

codesurvey/database.py
when-of-python/codesurvey
from .core import Source, Repo, RepoThunk, GitSource, GithubSampleSource, LocalSource, TestSource

__all__ = [
    'Source',
    'Repo',
    'RepoThunk',
    'GitSource',
    'GithubSampleSource',
    'LocalSource',
    'TestSource',
]

codesurvey/sources/__init__.py
when-of-python/codesurvey
"""Base classes and built-in Sources of code Repos to analyze."""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from functools import partial
import os
import os.path
from random import Random
import subprocess
from tempfile import mkdtemp, TemporaryDirectory
from typing import Any, Callable, Dict, Iterator, Mapping, Optional, Sequence, Union

import requests

from codesurvey.utils import logger, noop


class SourceError(Exception):
    """Raised for any failure of a Source to provide a Repo."""
    pass


@dataclass(frozen=True)
class Repo:
    """A repository of code that is accessible in a local directory in
    order to be analyzed."""

    source: 'Source'
    """Source the Repo is provided by."""

    key: str
    """Unique key of the Repo within its Source."""

    # TODO: Change to a pathlib.Path?
    path: str
    """Path to the local directory storing the Repo."""

    cleanup: Callable[[], None] = noop
    """Function to be called to remove or otherwise cleanup the Repo when
    analysis of it has finished."""

    metadata: Dict[str, Any] = field(default_factory=dict, compare=False)
    """Additional properties describing the Repo.

    The metadata structure may vary depending on the type of
    Source.

    """

    def __str__(self):
        return f'{self.source.name}:{self.key}'

    def __repr__(self):
        return f'{self.__class__.__name__}({self})'


@dataclass(frozen=True)
class RepoThunk:
    """An executable task to be run asynchronously to prepare a Repo."""

    source: 'Source'
    """Source the Repo is provided by."""

    key: str
    """Unique key of the Repo within its Source."""

    thunk: Callable[[], Repo]
    """Function to be called to prepare and return the Repo."""

    def __str__(self):
        return f'{self.source.name}:{self.key}'

    def __repr__(self):
        return f'{self.__class__.__name__}({self})'


class Source(ABC):
    """Provides Repos to be anaylyzed by CodeSurvey."""

    default_name: str
    """Default name to be assigned to Sources of this type if a custom
    name is not specified."""

    def __init__(self, *, name: Optional[str] = None):
        """
        Args:
            name: Name to identify the Source. If `None`, defaults to the
                Source type's default_name
        """
        self.name = self.default_name if name is None else name
        if self.name is None:
            raise ValueError('Analyzer name cannot be None')

    @abstractmethod
    def fetch_repo(self, repo_key: str) -> Repo:
        """Prepares the [Repo][codesurvey.sources.Repo] with the given
        `repo_key` for analysis.

        Typically called internally by repo_generator or by a
        RepoThunk, but also useful for inspecting a Repo given it's
        key from a survey result.

        """

    @abstractmethod
    def repo_generator(self) -> Iterator[Union[Repo, RepoThunk]]:
        """Generator yielding [Repos][codesurvey.sources.Repo] ready for
        analysis or [RepoThunks][codesurvey.sources.RepoThunk] that
        can be executed to prepare them for analysis."""

    def __str__(self):
        return self.name

    def __repr__(self):
        return f'{self.__class__.__name__}({self})'

    def repo(self, **kwargs):
        """Internal helper to generate a Repo for this Source. Takes the same
        arguments as Repo except for source."""
        return Repo(source=self, **kwargs)

    def repo_thunk(self, **kwargs):
        """Internal helper to generate a RepoThunk for this Source. Takes the
        same arguments as RepoThunk except for source."""
        return RepoThunk(source=self, **kwargs)


class LocalSource(Source):
    """
    Source of Repos from local filesystem directories.

    Example usage:

    ```python
    LocalSource([
        'path/to/my-source-code-directory',
    ])
    ```

    """
    default_name = 'local'

    def __init__(self, dirs: Sequence[str], *, name: Optional[str] = None):
        """
        Args:
            dirs: Paths to the local source code directory of each Repo
            name: Name to identify the Source. If `None`, defaults to 'local'.
        """
        self.dirs = dirs
        super().__init__(name=name)

    def fetch_repo(self, repo_key: str) -> Repo:
        return self.repo(
            key=repo_key,
            path=repo_key,
        )

    def repo_generator(self) -> Iterator[Repo]:
        for repo_dir in self.dirs:
            yield self.fetch_repo(repo_dir)


def fetch_git_repo(clone_url: str) -> str:
    """Helper function to clone a Git repository from the given URL."""
    temp_dir = mkdtemp()
    try:
        # Perform a shallow clone to reduce the download size.
        subprocess.run(['git', 'clone', '--depth', '1', clone_url, temp_dir],
                       capture_output=True, text=True, check=True)
    except subprocess.CalledProcessError as ex:
        raise SourceError((f'Failed to Git clone "{clone_url}" '
                           f'with error code {ex.returncode}\n'
                           f'> STDOUT {ex.stdout}\n'
                           f'> STDERR {ex.stderr}\n'))
    return temp_dir


class GitSource(Source):
    """
    Source of Repos from remote Git repositories.

    Repos are downloaded into a local directory for analysis.

    Example usage:

    ```python
    GitSource([
        'https://github.com/whenofpython/codesurvey',
    ])
    ```

    """
    default_name = 'git'

    def __init__(self, repo_urls: Sequence[str], *, name: Optional[str] = None):
        """
        Args:
            repo_urls: URLs of remote Git repositories.
            name: Name to identify the Source. If `None`, defaults to 'git'.
        """
        self.repo_urls = repo_urls
        super().__init__(name=name)

    def fetch_repo(self, repo_key: str) -> Repo:
        try:
            temp_dir = fetch_git_repo(repo_key)
        except SourceError as ex:
            raise SourceError(f'Source {self} failed to clone from GitHub: {ex}')
        return self.repo(
            key=repo_key,
            path=temp_dir,
            # When the repo is finished being used, the temp_dir
            # should be deleted:
            cleanup=partial(TemporaryDirectory._rmtree, temp_dir),  # type: ignore[attr-defined]
        )

    def repo_generator(self) -> Iterator[RepoThunk]:
        for repo_url in self.repo_urls:
            yield self.repo_thunk(
                key=repo_url,
                thunk=partial(self.fetch_repo, repo_url),
            )


class GithubSampleSource(Source):
    """Source of Repos sampled from GitHub's search API.

    Repos are sampled from randomly selected pages of GitHub search
    results, and downloaded to a temporary directory for analysis.

    For explanations of GitHub search parameters, see:
    https://docs.github.com/en/free-pro-team@latest/rest/search/search#search-repositories

    GitHub authentication credentials can be provided to increase rate
    limits. See:
    https://docs.github.com/en/rest/overview/authenticating-to-the-rest-api

    Example usage:

    ```python
    GithubSampleSource(language='python')
    ```

    """
    default_name = 'github_sample'

    REPOS_PER_PAGE = 100
    # GitHub only returns the first 1,000 search results
    MAX_RESULTS = 1000

    def __init__(self, *,
                 search_query: str = '',
                 language: Optional[str],
                 max_kb: Optional[int] = 50_000,
                 sort: str = 'updated',
                 auth_username: Optional[str] = None,
                 auth_token: Optional[str] = None,
                 random_seed: Optional[int] = None,
                 name: Optional[str] = None):
        """
        Args:
            search_query: An optional search query for GitHub search.
            language: An optional constraint for GitHub's repository
                language tag.
            max_kb: To avoid downloading excessively large repositories,
                limits the maximum kilobyte size of sampled Repos.
            sort: Sort order for GitHub search. Important as GitHub will
                only return the first 1000 pages of search results to sample
                from. Defaults to searching for recently updated repositories.
            auth_username: Username for GitHub authentication.
            auth_token: Token for GitHub authentication.
            random_seed: Random seed for sampling pages of search results.
                If `None`, a randomly selected seed is used.
            name: Name to identify the Source. If `None`, defaults
                to 'github_sample'.
        """
        self.search_query = search_query
        self.language = language
        self.max_kb = max_kb
        self.sort = sort
        self.auth = (auth_username, auth_token) if auth_username and auth_token else None
        self.random_seed = random_seed
        super().__init__(name=name)

    def _search_repos(self, *, page: int = 1) -> dict:
        """
        Makes a GitHub repo search API call for the specified page index.

        Returns a dictionary containing result metadata (total number of pages)
        and a list of dictionaries containing metadata for found repos.

        See:

        * https://docs.github.com/en/rest/search#search-repositories
        * https://docs.github.com/en/search-github/searching-on-github/searching-for-repositories
        """
        q_parts = []
        if self.search_query:
            q_parts.append(self.search_query)
        if self.language is not None:
            q_parts.append(f'language:{self.language}')
        if self.max_kb is not None:
            q_parts.append(f'size:<={self.max_kb}')

        params: Dict[str, Union[str, int]] = {
            'q': ' '.join(q_parts),
            'sort': self.sort,
            'per_page': self.REPOS_PER_PAGE,
            'page': page,
        }

        r = requests.get(
            'https://api.github.com/search/repositories',
            auth=self.auth,
            params=params,
        )
        r_json = r.json()
        return {
            # Return the total number of result pages that can be
            # sampled from.
            'page_count': max(self.MAX_RESULTS, r_json['total_count']) / self.REPOS_PER_PAGE,
            # Restrict the list of returned repos to those that
            # have the (optionally) specified language.
            'repos': [
                item for item in r_json['items']
                if (self.language is None or (str(item['language']).lower() == self.language))
            ],
        }

    def _clone_repo(self, repo_data: dict) -> Repo:
        """Helper to clone a Git repository given repo_data from the GitHub repos API."""
        try:
            temp_dir = fetch_git_repo(repo_data['clone_url'])
        except SourceError as ex:
            raise SourceError(f'Source {self} failed to clone from GitHub: {ex}')
        return self.repo(
            key=repo_data['full_name'],
            path=temp_dir,
            # When the repo is finished being used, the temp_dir
            # should be deleted:
            cleanup=partial(TemporaryDirectory._rmtree, temp_dir),  # type: ignore[attr-defined]
            metadata={
                'stars': repo_data['stargazers_count'],
            },
        )

    def fetch_repo(self, repo_key: str) -> Repo:
        r = requests.get(
            'https://api.github.com/repos/{repo_key}',
            auth=self.auth,
        )
        return self._clone_repo(r.json())

    def repo_generator(self) -> Iterator[RepoThunk]:
        rng = Random(self.random_seed)
        page_count = 1
        while True:
            logger.info(f'Source "{self}" searching GitHub for repos')
            search_result = self._search_repos(page=rng.randint(1, page_count))
            page_count = search_result['page_count']
            for repo_data in search_result['repos']:
                yield self.repo_thunk(
                    key=repo_data['full_name'],
                    thunk=partial(self._clone_repo, repo_data),
                )


class TestSource(Source):
    """Creates a single Repo in a temporary directory with specified files
    and contents.

    Only use with trusted paths, as paths are not checked for absolute
    or parent directory navigation.

    """
    default_name = 'test'

    def __init__(self, path_to_content: Mapping[str, str], *, name: Optional[str] = None):
        """
        Args:
            path_to_content: Mapping of paths to contents for files to create
                in a test Repo directory
            name: Name to identify the Source. If `None`, defaults to 'test'.
        """
        self.path_to_content = path_to_content
        super().__init__(name=name)

    def fetch_repo(self, repo_key: str) -> Repo:
        return self.repo(
            key=repo_key,
            path=repo_key,
            # When the repo is finished being used, the temporary
            # directory should be deleted:
            cleanup=partial(TemporaryDirectory._rmtree, repo_key),  # type: ignore[attr-defined]
        )

    def repo_generator(self) -> Iterator[Repo]:
        temp_dir = mkdtemp()
        for path, content in self.path_to_content.items():
            path_head, path_tail = os.path.split(path)
            path_dir = os.path.join(temp_dir, path_head)
            os.makedirs(path_dir, exist_ok=True)
            with open(os.path.join(path_dir, path_tail), 'w') as path_file:
                path_file.write(content)
        yield self.fetch_repo(temp_dir)

codesurvey/sources/core.py
when-of-python/codesurvey
"""Common utility functions."""

import logging
from typing import Dict, Hashable, List, Sequence, TypeVar


def get_logger():
    """
    Return a logger configured for use by codesurvey modules.
    """
    logger = logging.getLogger('codesurvey')
    logger_handler = logging.StreamHandler()
    logger.addHandler(logger_handler)
    logger_formatter = logging.Formatter('%(asctime)s - %(levelname)s: %(message)s',
                                         '%Y-%m-%d %H:%M:%S')
    logger_handler.setFormatter(logger_formatter)
    logger.setLevel(logging.INFO)
    return logger


logger = get_logger()
"""`logging.Logger` object that codesurvey logs events to during survey runs.

Can be used to customize logging:

```python
import logging
from codesurvey import logger

logger.setLevel(logging.ERROR)
```
"""


class BreakException(Exception):
    """
    Exception used for breaking out of nested loops.
    """


def noop(*args, **kwargs):
    """
    Function that will do nothing when called with any arguments.
    """
    pass


T = TypeVar('T', bound=Hashable)


def get_duplicates(items: Sequence[T]) -> List[T]:
    """
    Returns a list of any duplicate values in items.
    """
    seen = set()
    duplicates = []
    for item in items:
        if item in seen:
            if item not in duplicates:
                duplicates.append(item)
        else:
            seen.add(item)
    return duplicates


def recursive_update(dict_a: Dict, dict_b: Dict):
    """
    Recursively update dict_a by merging nested dictionaries from dict_b.
    """
    for key, b_value in dict_b.items():
        if key not in dict_a:
            dict_a[key] = b_value
        elif isinstance(b_value, dict):
            recursive_update(dict_a[key], b_value)

codesurvey/utils.py
when-of-python/codesurvey
from pprint import pprint

from codesurvey import CodeSurvey
from codesurvey.sources import GithubSampleSource
from codesurvey.analyzers.python import PythonAstAnalyzer
from codesurvey.analyzers.python.features import py_ast_feature_finder_with_transform


@py_ast_feature_finder_with_transform('add_note', xpath='Call/func/Attribute')
def has_add_note(func_attribute_el):
    if func_attribute_el.get('attr') == 'add_note':
        return func_attribute_el
    return None


survey = CodeSurvey(
    db_filepath='examples/add_note.sqlite3',
    sources=[
        GithubSampleSource(language='python'),
    ],
    analyzers=[
        PythonAstAnalyzer(
            feature_finders=[
                has_add_note,
            ],
        ),
    ],
    max_workers=3,
    use_saved_features=False,
)

# Runs until Ctrl-C
survey.run()

print('===== repo_features =====')
pprint(survey.get_repo_features()[:1])
print('===== code_features =====')
pprint(survey.get_code_features()[:1])
print('===== survey_tree =====')
pprint(survey.get_survey_tree(), depth=3)

examples/add_note.py
when-of-python/codesurvey
from pprint import pprint

from codesurvey import CodeSurvey
from codesurvey.sources import LocalSource, GithubSampleSource, GitSource
from codesurvey.analyzers.python import PythonAstAnalyzer
from codesurvey.analyzers.python.features import py_module_feature_finder, has_set


has_dataclasses = py_module_feature_finder('dataclasses', modules=['dataclasses'])

survey = CodeSurvey(
    db_filepath='examples/basic.sqlite3',
    sources=[
        LocalSource([
            '.',
        ]),
        GitSource([
            'https://github.com/when-of-python/blog.git',
        ]),
        GithubSampleSource(language='python'),
    ],
    analyzers=[
        PythonAstAnalyzer(
            feature_finders=[
                has_set,
                has_dataclasses,
            ],
        ),
    ],
    max_workers=3,
    use_saved_features=False,
)

survey.run(
    max_repos=4
)

print('===== repo_features =====')
pprint(survey.get_repo_features()[:1])
print('===== code_features =====')
pprint(survey.get_code_features()[:1])
print('===== survey_tree =====')
pprint(survey.get_survey_tree(), depth=3)

print('===== analyzer test =====')
analyzer = PythonAstAnalyzer(
    feature_finders=[
        has_set,
        has_dataclasses,
    ],
)
pprint(analyzer.test('''
import dataclasses
'''))

examples/basic.py
when-of-python/codesurvey
from codesurvey import __version__


def test_version():
    assert __version__ == '0.1.3'

tests/test_codesurvey.py
when-of-python/codesurvey
from codesurvey.utils import (
    get_duplicates,
    recursive_update,
)


def test_get_duplicates():
    assert get_duplicates([]) == []
    assert get_duplicates([3, 2, 1]) == []
    assert get_duplicates([1, 2, 3, 3, 1, 2, 1]) == [3, 1, 2]
    assert get_duplicates(['1', '2', '3', '3', '1', '2', '1']) == ['3', '1', '2']
    assert get_duplicates([None, None]) == [None]


def test_recursive_update():
    test_dict = {}
    recursive_update(test_dict, {})
    assert test_dict == {}

    test_dict = {'a': {1: 'I'}, 'b': {}}
    recursive_update(test_dict, {})
    assert test_dict == {'a': {1: 'I'}, 'b': {}}

    test_dict = {'a': {1: 'I'}, 'b': {}}
    recursive_update(test_dict, {'a': {1: 'i', 2: 'ii'}, 'b': 'test'})
    assert test_dict == {'a': {1: 'I', 2: 'ii'}, 'b': {}}

tests/test_utils.py
when-of-python/codesurvey