<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""File logging handler for tasks."""
import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Optional

import requests

from airflow.configuration import AirflowConfigException, conf
from airflow.utils.helpers import parse_template_string

if TYPE_CHECKING:
    from airflow.models import TaskInstance

class FileTaskHandler(logging.Handler):
    """
    FileTaskHandler is a python log handler that handles and reads
    task instance logs. It creates and delegates log handling
    to `logging.FileHandler` after receiving task instance context.
    It reads logs from task instance's host machine.

    :param base_log_folder: Base log folder to place logs.
    :param filename_template: template filename string
    """

    def __init__(self, base_log_folder: str, filename_template: str):
        super().__init__()
        self.handler = None  # type: Optional[logging.FileHandler]
        self.local_base = base_log_folder
        self.filename_template, self.filename_jinja_template = parse_template_string(filename_template)

    def set_context(self, ti: "TaskInstance"):
        """
        Provide task_instance context to airflow task handler.

        :param ti: task instance object
        """
        local_loc = self._init_file(ti)
        self.handler = logging.FileHandler(local_loc, encoding='utf-8')
        if self.formatter:
            self.handler.setFormatter(self.formatter)
        self.handler.setLevel(self.level)

    def emit(self, record):
        if self.handler:
            self.handler.emit(record)

    def flush(self):
        if self.handler:
            self.handler.flush()

    def close(self):
        if self.handler:
            self.handler.close()

    def _render_filename(self, ti, try_number):
        if self.filename_jinja_template:
            if hasattr(ti, 'task'):
                jinja_context = ti.get_template_context()
                jinja_context['try_number'] = try_number
            else:
                jinja_context = {
                    'ti': ti,
                    'ts': ti.execution_date.isoformat(),
                    'try_number': try_number,
                }
            return self.filename_jinja_template.render(**jinja_context)

        return self.filename_template.format(
            dag_id=ti.dag_id,
            task_id=ti.task_id,
            execution_date=ti.execution_date.isoformat(),
            try_number=try_number,
        )

    def _read_grouped_logs(self):
        return False

    def _read(self, ti, try_number, metadata=None):  # pylint: disable=unused-argument
        """
        Template method that contains custom logic of reading
        logs given the try_number.

        :param ti: task instance record
        :param try_number: current try_number to read log from
        :param metadata: log metadata,
                         can be used for steaming log reading and auto-tailing.
        :return: log message as a string and metadata.
        """
        # Task instance here might be different from task instance when
        # initializing the handler. Thus explicitly getting log location
        # is needed to get correct log path.
        log_relative_path = self._render_filename(ti, try_number)
        location = os.path.join(self.local_base, log_relative_path)

        log = ""

        if os.path.exists(location):
            try:
                with open(location) as file:
                    log += f"*** Reading local file: {location}\n"
                    log += "".join(file.readlines())
            except Exception as e:  # pylint: disable=broad-except
                log = f"*** Failed to load local log file: {location}\n"
                log += f"*** {str(e)}\n"
        elif conf.get('core', 'executor') == 'KubernetesExecutor':  # pylint: disable=too-many-nested-blocks
            try:
                from airflow.kubernetes.kube_client import get_kube_client

                kube_client = get_kube_client()

                if len(ti.hostname) >= 63:
                    # Kubernetes takes the pod name and truncates it for the hostname. This truncated hostname
                    # is returned for the fqdn to comply with the 63 character limit imposed by DNS standards
                    # on any label of a FQDN.
                    pod_list = kube_client.list_namespaced_pod(conf.get('kubernetes', 'namespace'))
                    matches = [
                        pod.metadata.name
                        for pod in pod_list.items
                        if pod.metadata.name.startswith(ti.hostname)
                    ]
                    if len(matches) == 1:
                        if len(matches[0]) > len(ti.hostname):
                            ti.hostname = matches[0]

                log += '*** Trying to get logs (last 100 lines) from worker pod {} ***\n\n'.format(
                    ti.hostname
                )

                res = kube_client.read_namespaced_pod_log(
                    name=ti.hostname,
                    namespace=conf.get('kubernetes', 'namespace'),
                    container='base',
                    follow=False,
                    tail_lines=100,
                    _preload_content=False,
                )

                for line in res:
                    log += line.decode()

            except Exception as f:  # pylint: disable=broad-except
                log += f'*** Unable to fetch logs from worker pod {ti.hostname} ***\n{str(f)}\n\n'
        else:
            url = os.path.join("http://{ti.hostname}:{worker_log_server_port}/log", log_relative_path).format(
                ti=ti, worker_log_server_port=conf.get('celery', 'WORKER_LOG_SERVER_PORT')
            )
            log += f"*** Log file does not exist: {location}\n"
            log += f"*** Fetching from: {url}\n"
            try:
                timeout = None  # No timeout
                try:
                    timeout = conf.getint('webserver', 'log_fetch_timeout_sec')
                except (AirflowConfigException, ValueError):
                    pass

                response = requests.get(url, timeout=timeout)
                response.encoding = "utf-8"

                # Check if the resource was properly fetched
                response.raise_for_status()

                log += '\n' + response.text
            except Exception as e:  # pylint: disable=broad-except
                log += f"*** Failed to fetch log file from worker. {str(e)}\n"

        return log, {'end_of_log': True}

    def read(self, task_instance, try_number=None, metadata=None):
        """
        Read logs of given task instance from local machine.

        :param task_instance: task instance object
        :param try_number: task instance try_number to read logs from. If None
                           it returns all logs separated by try_number
        :param metadata: log metadata,
                         can be used for steaming log reading and auto-tailing.
        :return: a list of listed tuples which order log string by host
        """
        # Task instance increments its try number when it starts to run.
        # So the log for a particular task try will only show up when
        # try number gets incremented in DB, i.e logs produced the time
        # after cli run and before try_number + 1 in DB will not be displayed.

        if try_number is None:
            next_try = task_instance.next_try_number
            try_numbers = list(range(1, next_try))
        elif try_number < 1:
            logs = [
                [('default_host', f'Error fetching the logs. Try number {try_number} is invalid.')],
            ]
            return logs, [{'end_of_log': True}]
        else:
            try_numbers = [try_number]

        logs = [''] * len(try_numbers)
        metadata_array = [{}] * len(try_numbers)
        for i, try_number_element in enumerate(try_numbers):
            log, metadata = self._read(task_instance, try_number_element, metadata)
            # es_task_handler return logs grouped by host. wrap other handler returning log string
            # with default/ empty host so that UI can render the response in the same way
            logs[i] = log if self._read_grouped_logs() else [(task_instance.hostname, log)]
            metadata_array[i] = metadata

        return logs, metadata_array

    def _init_file(self, ti):
        """
        Create log directory and give it correct permissions.

        :param ti: task instance object
        :return: relative log path of the given task instance
        """
        # To handle log writing when tasks are impersonated, the log files need to
        # be writable by the user that runs the Airflow command and the user
        # that is impersonated. This is mainly to handle corner cases with the
        # SubDagOperator. When the SubDagOperator is run, all of the operators
        # run under the impersonated user and create appropriate log files
        # as the impersonated user. However, if the user manually runs tasks
        # of the SubDagOperator through the UI, then the log files are created
        # by the user that runs the Airflow command. For example, the Airflow
        # run command may be run by the `airflow_sudoable` user, but the Airflow
        # tasks may be run by the `airflow` user. If the log files are not
        # writable by both users, then it's possible that re-running a task
        # via the UI (or vice versa) results in a permission error as the task
        # tries to write to a log file created by the other user.
        relative_path = self._render_filename(ti, ti.try_number)
        full_path = os.path.join(self.local_base, relative_path)
        directory = os.path.dirname(full_path)
        # Create the log file and give it group writable permissions
        # TODO(aoen): Make log dirs and logs globally readable for now since the SubDag
        # operator is not compatible with impersonation (e.g. if a Celery executor is used
        # for a SubDag operator and the SubDag operator has a different owner than the
        # parent DAG)
        Path(directory).mkdir(mode=0o777, parents=True, exist_ok=True)

        if not os.path.exists(full_path):
            open(full_path, "a").close()
            # TODO: Investigate using 444 instead of 666.
            os.chmod(full_path, 0o666)

        return full_path
<commit_msg>Fix permission error on non-POSIX filesystem (#13121)<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""File logging handler for tasks."""
import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Optional

import requests

from airflow.configuration import AirflowConfigException, conf
from airflow.utils.helpers import parse_template_string

if TYPE_CHECKING:
    from airflow.models import TaskInstance

class FileTaskHandler(logging.Handler):
    """
    FileTaskHandler is a python log handler that handles and reads
    task instance logs. It creates and delegates log handling
    to `logging.FileHandler` after receiving task instance context.
    It reads logs from task instance's host machine.

    :param base_log_folder: Base log folder to place logs.
    :param filename_template: template filename string
    """

    def __init__(self, base_log_folder: str, filename_template: str):
        super().__init__()
        self.handler = None  # type: Optional[logging.FileHandler]
        self.local_base = base_log_folder
        self.filename_template, self.filename_jinja_template = parse_template_string(filename_template)

    def set_context(self, ti: "TaskInstance"):
        """
        Provide task_instance context to airflow task handler.

        :param ti: task instance object
        """
        local_loc = self._init_file(ti)
        self.handler = logging.FileHandler(local_loc, encoding='utf-8')
        if self.formatter:
            self.handler.setFormatter(self.formatter)
        self.handler.setLevel(self.level)

    def emit(self, record):
        if self.handler:
            self.handler.emit(record)

    def flush(self):
        if self.handler:
            self.handler.flush()

    def close(self):
        if self.handler:
            self.handler.close()

    def _render_filename(self, ti, try_number):
        if self.filename_jinja_template:
            if hasattr(ti, 'task'):
                jinja_context = ti.get_template_context()
                jinja_context['try_number'] = try_number
            else:
                jinja_context = {
                    'ti': ti,
                    'ts': ti.execution_date.isoformat(),
                    'try_number': try_number,
                }
            return self.filename_jinja_template.render(**jinja_context)

        return self.filename_template.format(
            dag_id=ti.dag_id,
            task_id=ti.task_id,
            execution_date=ti.execution_date.isoformat(),
            try_number=try_number,
        )

    def _read_grouped_logs(self):
        return False

    def _read(self, ti, try_number, metadata=None):  # pylint: disable=unused-argument
        """
        Template method that contains custom logic of reading
        logs given the try_number.

        :param ti: task instance record
        :param try_number: current try_number to read log from
        :param metadata: log metadata,
                         can be used for steaming log reading and auto-tailing.
        :return: log message as a string and metadata.
        """
        # Task instance here might be different from task instance when
        # initializing the handler. Thus explicitly getting log location
        # is needed to get correct log path.
        log_relative_path = self._render_filename(ti, try_number)
        location = os.path.join(self.local_base, log_relative_path)

        log = ""

        if os.path.exists(location):
            try:
                with open(location) as file:
                    log += f"*** Reading local file: {location}\n"
                    log += "".join(file.readlines())
            except Exception as e:  # pylint: disable=broad-except
                log = f"*** Failed to load local log file: {location}\n"
                log += f"*** {str(e)}\n"
        elif conf.get('core', 'executor') == 'KubernetesExecutor':  # pylint: disable=too-many-nested-blocks
            try:
                from airflow.kubernetes.kube_client import get_kube_client

                kube_client = get_kube_client()

                if len(ti.hostname) >= 63:
                    # Kubernetes takes the pod name and truncates it for the hostname. This truncated hostname
                    # is returned for the fqdn to comply with the 63 character limit imposed by DNS standards
                    # on any label of a FQDN.
                    pod_list = kube_client.list_namespaced_pod(conf.get('kubernetes', 'namespace'))
                    matches = [
                        pod.metadata.name
                        for pod in pod_list.items
                        if pod.metadata.name.startswith(ti.hostname)
                    ]
                    if len(matches) == 1:
                        if len(matches[0]) > len(ti.hostname):
                            ti.hostname = matches[0]

                log += '*** Trying to get logs (last 100 lines) from worker pod {} ***\n\n'.format(
                    ti.hostname
                )

                res = kube_client.read_namespaced_pod_log(
                    name=ti.hostname,
                    namespace=conf.get('kubernetes', 'namespace'),
                    container='base',
                    follow=False,
                    tail_lines=100,
                    _preload_content=False,
                )

                for line in res:
                    log += line.decode()

            except Exception as f:  # pylint: disable=broad-except
                log += f'*** Unable to fetch logs from worker pod {ti.hostname} ***\n{str(f)}\n\n'
        else:
            url = os.path.join("http://{ti.hostname}:{worker_log_server_port}/log", log_relative_path).format(
                ti=ti, worker_log_server_port=conf.get('celery', 'WORKER_LOG_SERVER_PORT')
            )
            log += f"*** Log file does not exist: {location}\n"
            log += f"*** Fetching from: {url}\n"
            try:
                timeout = None  # No timeout
                try:
                    timeout = conf.getint('webserver', 'log_fetch_timeout_sec')
                except (AirflowConfigException, ValueError):
                    pass

                response = requests.get(url, timeout=timeout)
                response.encoding = "utf-8"

                # Check if the resource was properly fetched
                response.raise_for_status()

                log += '\n' + response.text
            except Exception as e:  # pylint: disable=broad-except
                log += f"*** Failed to fetch log file from worker. {str(e)}\n"

        return log, {'end_of_log': True}

    def read(self, task_instance, try_number=None, metadata=None):
        """
        Read logs of given task instance from local machine.

        :param task_instance: task instance object
        :param try_number: task instance try_number to read logs from. If None
                           it returns all logs separated by try_number
        :param metadata: log metadata,
                         can be used for steaming log reading and auto-tailing.
        :return: a list of listed tuples which order log string by host
        """
        # Task instance increments its try number when it starts to run.
        # So the log for a particular task try will only show up when
        # try number gets incremented in DB, i.e logs produced the time
        # after cli run and before try_number + 1 in DB will not be displayed.

        if try_number is None:
            next_try = task_instance.next_try_number
            try_numbers = list(range(1, next_try))
        elif try_number < 1:
            logs = [
                [('default_host', f'Error fetching the logs. Try number {try_number} is invalid.')],
            ]
            return logs, [{'end_of_log': True}]
        else:
            try_numbers = [try_number]

        logs = [''] * len(try_numbers)
        metadata_array = [{}] * len(try_numbers)
        for i, try_number_element in enumerate(try_numbers):
            log, metadata = self._read(task_instance, try_number_element, metadata)
            # es_task_handler return logs grouped by host. wrap other handler returning log string
            # with default/ empty host so that UI can render the response in the same way
            logs[i] = log if self._read_grouped_logs() else [(task_instance.hostname, log)]
            metadata_array[i] = metadata

        return logs, metadata_array

    def _init_file(self, ti):
        """
        Create log directory and give it correct permissions.

        :param ti: task instance object
        :return: relative log path of the given task instance
        """
        # To handle log writing when tasks are impersonated, the log files need to
        # be writable by the user that runs the Airflow command and the user
        # that is impersonated. This is mainly to handle corner cases with the
        # SubDagOperator. When the SubDagOperator is run, all of the operators
        # run under the impersonated user and create appropriate log files
        # as the impersonated user. However, if the user manually runs tasks
        # of the SubDagOperator through the UI, then the log files are created
        # by the user that runs the Airflow command. For example, the Airflow
        # run command may be run by the `airflow_sudoable` user, but the Airflow
        # tasks may be run by the `airflow` user. If the log files are not
        # writable by both users, then it's possible that re-running a task
        # via the UI (or vice versa) results in a permission error as the task
        # tries to write to a log file created by the other user.
        relative_path = self._render_filename(ti, ti.try_number)
        full_path = os.path.join(self.local_base, relative_path)
        directory = os.path.dirname(full_path)
        # Create the log file and give it group writable permissions
        # TODO(aoen): Make log dirs and logs globally readable for now since the SubDag
        # operator is not compatible with impersonation (e.g. if a Celery executor is used
        # for a SubDag operator and the SubDag operator has a different owner than the
        # parent DAG)
        Path(directory).mkdir(mode=0o777, parents=True, exist_ok=True)

        if not os.path.exists(full_path):
            open(full_path, "a").close()
            # TODO: Investigate using 444 instead of 666.
            try:
                os.chmod(full_path, 0o666)
            except OSError:
                logging.warning("OSError while change ownership of the log file")

        return full_path
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import unicode_literals

import argparse
import ast
import collections
import io
import sys

from tokenize_rt import ESCAPED_NL
from tokenize_rt import src_to_tokens
from tokenize_rt import Token
from tokenize_rt import tokens_to_src
from tokenize_rt import UNIMPORTANT_WS

Offset = collections.namedtuple('Offset', ('line', 'utf8_byte_offset'))
Call = collections.namedtuple('Call', ('node', 'star_args', 'arg_offsets'))
Func = collections.namedtuple('Func', ('node', 'star_args', 'arg_offsets'))
Literal = collections.namedtuple('Literal', ('node', 'backtrack'))
Literal.__new__.__defaults__ = (False,)
Fix = collections.namedtuple('Fix', ('braces', 'multi_arg', 'initial_indent'))

NEWLINES = frozenset((ESCAPED_NL, 'NEWLINE', 'NL'))
NON_CODING_TOKENS = frozenset(('COMMENT', ESCAPED_NL, 'NL', UNIMPORTANT_WS))
INDENT_TOKENS = frozenset(('INDENT', UNIMPORTANT_WS))
START_BRACES = frozenset(('(', '{', '['))
END_BRACES = frozenset((')', '}', ']'))

def ast_parse(contents_text):
    return ast.parse(contents_text.encode('UTF-8'))

def _to_offset(node):
    candidates = [node]
    while candidates:
        candidate = candidates.pop()
        if hasattr(candidate, 'lineno'):
            return Offset(candidate.lineno, candidate.col_offset)
        elif hasattr(candidate, '_fields'):  # pragma: no cover (PY35+)
            for field in reversed(candidate._fields):
                candidates.append(getattr(candidate, field))
    else:
        raise AssertionError(node)

if sys.version_info < (3, 5):  # pragma: no cover (<PY35)
    def _is_star_arg(node):
        return False
else:  # pragma: no cover (PY35+)
    def _is_star_arg(node):
        return isinstance(node, ast.Starred)

def _is_star_star_kwarg(node):
    return isinstance(node, ast.keyword) and node.arg is None

class FindNodes(ast.NodeVisitor):
    def __init__(self):
        self.calls = {}
        self.funcs = {}
        self.literals = {}
        self.has_new_syntax = False

    def _visit_literal(self, node, key='elts', is_multiline=False, **kwargs):
        orig = node.lineno

        for elt in getattr(node, key):
            if elt.lineno > orig:
                is_multiline = True
            if _is_star_arg(elt):  # pragma: no cover (PY35+)
                self.has_new_syntax = True

        if is_multiline:
            key = Offset(node.lineno, node.col_offset)
            self.literals[key] = Literal(node, **kwargs)
            self.generic_visit(node)

    visit_Set = visit_List = _visit_literal

    def visit_Dict(self, node):
        # unpackings are represented as a `None` key
        if None in node.keys:  # pragma: no cover (PY35+)
            self.has_new_syntax = True
        self._visit_literal(node, key='values')

    def visit_Tuple(self, node):
        # tuples lie about things, so we pretend they are all multiline
        # and tell the later machinery to backtrack
        self._visit_literal(node, is_multiline=True, backtrack=True)

    def visit_Call(self, node):
        orig = node.lineno

        argnodes = node.args + node.keywords
        py2_starargs = getattr(node, 'starargs', None)
        if py2_starargs:  # pragma: no cover (<PY35)
            argnodes.append(py2_starargs)
        py2_kwargs = getattr(node, 'kwargs', None)
        if py2_kwargs:  # pragma: no cover (<PY35)
            argnodes.append(py2_kwargs)

        arg_offsets = set()
        is_multiline = False
        has_starargs = bool(py2_starargs or py2_kwargs)
        for argnode in argnodes:
            if (
                    _is_star_arg(argnode) or
                    _is_star_star_kwarg(argnode)
            ):  # pragma: no cover (PY35+)
                has_starargs = True

            offset = _to_offset(argnode)
            # multiline strings have invalid position, ignore them
            if offset.utf8_byte_offset != -1:  # pragma: no branch (cpy bug)
                if offset.line > orig:
                    is_multiline = True
                arg_offsets.add(offset)

        # If the sole argument is a generator, don't add a trailing comma as
        # this breaks lib2to3 based tools
        only_a_generator = (
            len(argnodes) == 1 and isinstance(argnodes[0], ast.GeneratorExp)
        )

        if is_multiline and not only_a_generator:
            key = Offset(node.lineno, node.col_offset)
            self.calls[key] = Call(node, has_starargs, arg_offsets)

        if (
                sum(_is_star_arg(n) for n in node.args) > 1 or
                sum(_is_star_star_kwarg(n) for n in node.keywords) > 1
        ):  # pragma: no cover (PY35+)
            self.has_new_syntax = True

        self.generic_visit(node)

    def visit_FunctionDef(self, node):
        has_starargs = False
        args = list(node.args.args)

        if node.args.vararg:
            if isinstance(node.args.vararg, ast.AST):  # pragma: no cover (py3)
                args.append(node.args.vararg)
            has_starargs = True
        if node.args.kwarg:
            if isinstance(node.args.kwarg, ast.AST):  # pragma: no cover (py3)
                args.append(node.args.kwarg)
            has_starargs = True
        py3_kwonlyargs = getattr(node.args, 'kwonlyargs', None)
        if py3_kwonlyargs:  # pragma: no cover (py3)
            args.extend(py3_kwonlyargs)
            has_starargs = True

        orig = node.lineno
        is_multiline = False
        offsets = set()
        for argnode in args:
            offset = _to_offset(argnode)
            if offset.line > orig:
                is_multiline = True
            offsets.add(offset)

        if is_multiline:
            key = Offset(node.lineno, node.col_offset)
            self.funcs[key] = Func(node, has_starargs, offsets)

        self.generic_visit(node)

def _find_simple(first_brace, tokens):
    brace_stack = [first_brace]
    multi_arg = False

    for i in range(first_brace + 1, len(tokens)):
        token = tokens[i]
        if token.src in START_BRACES:
            brace_stack.append(i)
        elif token.src in END_BRACES:
            brace_stack.pop()

        if len(brace_stack) == 1 and token.src == ',':
            multi_arg = True

        if not brace_stack:
            break
    else:
        raise AssertionError('Past end?')

    last_brace = i

    # This was not actually a multi-line call, despite the ast telling us that
    if tokens[first_brace].line == tokens[last_brace].line:
        return

    # determine the initial indentation
    i = first_brace
    while i >= 0 and tokens[i].name not in NEWLINES:
        i -= 1

    if i >= 0 and tokens[i + 1].name in INDENT_TOKENS:
        initial_indent = len(tokens[i + 1].src)
    else:
        initial_indent = 0

    return Fix((first_brace, last_brace), multi_arg, initial_indent)

def _find_call(call, i, tokens):
    # When we get a `call` object, the ast refers to it as this:
    #
    #     func_name(arg, arg, arg)
    #     ^ where ast points
    #
    # We care about the closing paren, in order to find it, we first walk
    # until we find an argument.  When we find an argument, we know the outer
    # paren we find is the function call paren
    #
    #     func_name(arg, arg, arg)
    #              ^ outer paren
    first_brace = None
    paren_stack = []
    for i in range(i, len(tokens)):
        token = tokens[i]
        if token.src == '(':
            paren_stack.append(i)
        # the ast lies to us about the beginning of parenthesized functions.
        # See #3. (why we make sure there's something to pop here)
        elif token.src == ')' and paren_stack:
            paren_stack.pop()

        if (token.line, token.utf8_byte_offset) in call.arg_offsets:
            first_brace = paren_stack[0]
            break
    else:
        raise AssertionError('Past end?')

    return _find_simple(first_brace, tokens)

def _find_literal(literal, i, tokens):
    # tuples are evil, we need to backtrack to find the opening paren
    if literal.backtrack:
        i -= 1
        while tokens[i].name in NON_CODING_TOKENS:
            i -= 1
        # Sometimes tuples don't even have a paren!
        # x = 1, 2, 3
        if tokens[i].src != '(':
            return

    return _find_simple(i, tokens)

def _fix_comma_and_unhug(fix_data, add_comma, tokens):
    first_brace, last_brace = fix_data.braces

    # Figure out if either of the braces are "hugging"
    hug_open = tokens[first_brace + 1].name not in NON_CODING_TOKENS
    hug_close = tokens[last_brace - 1].name not in NON_CODING_TOKENS
    if (
            not fix_data.multi_arg and
            hug_open and
            tokens[last_brace - 1].src in END_BRACES
    ):
        hug_open = hug_close = False

    # fix open hugging
    if hug_open:
        new_indent = fix_data.initial_indent + 4

        tokens[first_brace + 1:first_brace + 1] = [
            Token('NL', '\n'), Token(UNIMPORTANT_WS, ' ' * new_indent),
        ]
        last_brace += 2

        # Adust indentation for the rest of the things
        min_indent = None
        indents = []
        insert_indents = []
        for i in range(first_brace + 3, last_brace):
            if tokens[i - 1].name == 'NL':
                if tokens[i].name != UNIMPORTANT_WS:
                    min_indent = 0
                    insert_indents.append(i)
                else:
                    if min_indent is None:
                        min_indent = len(tokens[i].src)
                    elif len(tokens[i].src) < min_indent:
                        min_indent = len(tokens[i].src)
                    indents.append(i)

        for i in indents:
            oldlen = len(tokens[i].src)
            newlen = oldlen - min_indent + new_indent
            tokens[i] = tokens[i]._replace(src=' ' * newlen)
        for i in reversed(insert_indents):
            tokens.insert(i, Token(UNIMPORTANT_WS, ' ' * new_indent))
            last_brace += 1

    # fix close hugging
    if hug_close:
        tokens[last_brace:last_brace] = [
            Token('NL', '\n'),
            Token(UNIMPORTANT_WS, ' ' * fix_data.initial_indent),
        ]
        last_brace += 2

    # From there, we can walk backwards and decide whether a comma is needed
    i = last_brace - 1
    while tokens[i].name in NON_CODING_TOKENS:
        i -= 1

    # If we're not a hugging paren, we can insert a comma
    if add_comma and tokens[i].src != ',' and i + 1 != last_brace:
        tokens.insert(i + 1, Token('OP', ','))

def _fix_trailing_brace(fix_data, tokens):
    _, last_brace = fix_data.braces

    back_1 = tokens[last_brace - 1]
    back_2 = tokens[last_brace - 2]

    if (
            back_1.name == UNIMPORTANT_WS and
            back_2.name == 'NL' and
            len(back_1.src) != fix_data.initial_indent
    ):
        new_indent = fix_data.initial_indent * ' '
        tokens[last_brace - 1] = back_1._replace(src=new_indent)

def _changing_list(lst):
    i = 0
    while i < len(lst):
        yield i, lst[i]
        i += 1

def _fix_src(contents_text, py35_plus):
    try:
        ast_obj = ast_parse(contents_text)
    except SyntaxError:
        return contents_text

    visitor = FindNodes()
    visitor.visit(ast_obj)
    py35_plus = py35_plus or visitor.has_new_syntax

    tokens = src_to_tokens(contents_text)
    for i, token in _changing_list(tokens):
        key = Offset(token.line, token.utf8_byte_offset)
        add_comma = True
        fix_data = None

        if key in visitor.calls:
            call = visitor.calls[key]
            # Only fix stararg calls if asked to
            add_comma = not call.star_args or py35_plus
            fix_data = _find_call(call, i, tokens)
        elif key in visitor.funcs:
            func = visitor.funcs[key]
            # any amount of starargs excludes adding a comma for defs
            add_comma = not func.star_args
            # functions can be treated as calls
            fix_data = _find_call(func, i, tokens)
        # Handle parenthesized things
        elif token.src == '(':
            fix_data = _find_simple(i, tokens)
            add_comma = False

        if fix_data is not None:
            _fix_comma_and_unhug(fix_data, add_comma, tokens)

        # need to additionally handle literals afterwards as tuples report
        # their starting index as the first element, which may be one of the
        # above things.
        if key in visitor.literals:
            fix_data = _find_literal(visitor.literals[key], i, tokens)
            if fix_data is not None:
                _fix_comma_and_unhug(fix_data, True, tokens)

    # Need a second pass to fix trailing braces after indentation is fixed
    for i, token in _changing_list(tokens):
        key = Offset(token.line, token.utf8_byte_offset)
        fix_data = None

        if key in visitor.calls:
            fix_data = _find_call(visitor.calls[key], i, tokens)
        elif key in visitor.funcs:
            fix_data = _find_call(visitor.funcs[key], i, tokens)
        # Handle parenthesized things
        elif token.src == '(':
            fix_data = _find_simple(i, tokens)

        if fix_data is not None:
            _fix_trailing_brace(fix_data, tokens)

        # need to additionally handle literals afterwards as tuples report
        # their starting index as the first element, which may be one of the
        # above things.
        if key in visitor.literals:
            fix_data = _find_literal(visitor.literals[key], i, tokens)
            if fix_data is not None:
                _fix_trailing_brace(fix_data, tokens)

    return tokens_to_src(tokens)

def fix_file(filename, args):
    with open(filename, 'rb') as f:
        contents_bytes = f.read()

    try:
        contents_text_orig = contents_text = contents_bytes.decode('UTF-8')
    except UnicodeDecodeError:
        print('{} is non-utf-8 (not supported)'.format(filename))
        return 1

    contents_text = _fix_src(contents_text, args.py35_plus)

    if contents_text != contents_text_orig:
        print('Rewriting {}'.format(filename))
        with io.open(filename, 'w', encoding='UTF-8') as f:
            f.write(contents_text)
        return 1

    return 0

def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('filenames', nargs='*')
    parser.add_argument('--py35-plus', action='store_true')
    args = parser.parse_args(argv)

    ret = 0
    for filename in args.filenames:
        ret |= fix_file(filename, args)
    return ret

if __name__ == '__main__':
    exit(main())
<commit_msg>Change back to single pass now that it is top-down<commit_after>from __future__ import absolute_import
from __future__ import unicode_literals

import argparse
import ast
import collections
import io
import sys

from tokenize_rt import ESCAPED_NL
from tokenize_rt import src_to_tokens
from tokenize_rt import Token
from tokenize_rt import tokens_to_src
from tokenize_rt import UNIMPORTANT_WS

Offset = collections.namedtuple('Offset', ('line', 'utf8_byte_offset'))
Call = collections.namedtuple('Call', ('node', 'star_args', 'arg_offsets'))
Func = collections.namedtuple('Func', ('node', 'star_args', 'arg_offsets'))
Literal = collections.namedtuple('Literal', ('node', 'backtrack'))
Literal.__new__.__defaults__ = (False,)
Fix = collections.namedtuple('Fix', ('braces', 'multi_arg', 'initial_indent'))

NEWLINES = frozenset((ESCAPED_NL, 'NEWLINE', 'NL'))
NON_CODING_TOKENS = frozenset(('COMMENT', ESCAPED_NL, 'NL', UNIMPORTANT_WS))
INDENT_TOKENS = frozenset(('INDENT', UNIMPORTANT_WS))
START_BRACES = frozenset(('(', '{', '['))
END_BRACES = frozenset((')', '}', ']'))

def ast_parse(contents_text):
    return ast.parse(contents_text.encode('UTF-8'))

def _to_offset(node):
    candidates = [node]
    while candidates:
        candidate = candidates.pop()
        if hasattr(candidate, 'lineno'):
            return Offset(candidate.lineno, candidate.col_offset)
        elif hasattr(candidate, '_fields'):  # pragma: no cover (PY35+)
            for field in reversed(candidate._fields):
                candidates.append(getattr(candidate, field))
    else:
        raise AssertionError(node)

if sys.version_info < (3, 5):  # pragma: no cover (<PY35)
    def _is_star_arg(node):
        return False
else:  # pragma: no cover (PY35+)
    def _is_star_arg(node):
        return isinstance(node, ast.Starred)

def _is_star_star_kwarg(node):
    return isinstance(node, ast.keyword) and node.arg is None

class FindNodes(ast.NodeVisitor):
    def __init__(self):
        self.calls = {}
        self.funcs = {}
        self.literals = {}
        self.has_new_syntax = False

    def _visit_literal(self, node, key='elts', is_multiline=False, **kwargs):
        orig = node.lineno

        for elt in getattr(node, key):
            if elt.lineno > orig:
                is_multiline = True
            if _is_star_arg(elt):  # pragma: no cover (PY35+)
                self.has_new_syntax = True

        if is_multiline:
            key = Offset(node.lineno, node.col_offset)
            self.literals[key] = Literal(node, **kwargs)
            self.generic_visit(node)

    visit_Set = visit_List = _visit_literal

    def visit_Dict(self, node):
        # unpackings are represented as a `None` key
        if None in node.keys:  # pragma: no cover (PY35+)
            self.has_new_syntax = True
        self._visit_literal(node, key='values')

    def visit_Tuple(self, node):
        # tuples lie about things, so we pretend they are all multiline
        # and tell the later machinery to backtrack
        self._visit_literal(node, is_multiline=True, backtrack=True)

    def visit_Call(self, node):
        orig = node.lineno

        argnodes = node.args + node.keywords
        py2_starargs = getattr(node, 'starargs', None)
        if py2_starargs:  # pragma: no cover (<PY35)
            argnodes.append(py2_starargs)
        py2_kwargs = getattr(node, 'kwargs', None)
        if py2_kwargs:  # pragma: no cover (<PY35)
            argnodes.append(py2_kwargs)

        arg_offsets = set()
        is_multiline = False
        has_starargs = bool(py2_starargs or py2_kwargs)
        for argnode in argnodes:
            if (
                    _is_star_arg(argnode) or
                    _is_star_star_kwarg(argnode)
            ):  # pragma: no cover (PY35+)
                has_starargs = True

            offset = _to_offset(argnode)
            # multiline strings have invalid position, ignore them
            if offset.utf8_byte_offset != -1:  # pragma: no branch (cpy bug)
                if offset.line > orig:
                    is_multiline = True
                arg_offsets.add(offset)

        # If the sole argument is a generator, don't add a trailing comma as
        # this breaks lib2to3 based tools
        only_a_generator = (
            len(argnodes) == 1 and isinstance(argnodes[0], ast.GeneratorExp)
        )

        if is_multiline and not only_a_generator:
            key = Offset(node.lineno, node.col_offset)
            self.calls[key] = Call(node, has_starargs, arg_offsets)

        if (
                sum(_is_star_arg(n) for n in node.args) > 1 or
                sum(_is_star_star_kwarg(n) for n in node.keywords) > 1
        ):  # pragma: no cover (PY35+)
            self.has_new_syntax = True

        self.generic_visit(node)

    def visit_FunctionDef(self, node):
        has_starargs = False
        args = list(node.args.args)

        if node.args.vararg:
            if isinstance(node.args.vararg, ast.AST):  # pragma: no cover (py3)
                args.append(node.args.vararg)
            has_starargs = True
        if node.args.kwarg:
            if isinstance(node.args.kwarg, ast.AST):  # pragma: no cover (py3)
                args.append(node.args.kwarg)
            has_starargs = True
        py3_kwonlyargs = getattr(node.args, 'kwonlyargs', None)
        if py3_kwonlyargs:  # pragma: no cover (py3)
            args.extend(py3_kwonlyargs)
            has_starargs = True

        orig = node.lineno
        is_multiline = False
        offsets = set()
        for argnode in args:
            offset = _to_offset(argnode)
            if offset.line > orig:
                is_multiline = True
            offsets.add(offset)

        if is_multiline:
            key = Offset(node.lineno, node.col_offset)
            self.funcs[key] = Func(node, has_starargs, offsets)

        self.generic_visit(node)

def _find_simple(first_brace, tokens):
    brace_stack = [first_brace]
    multi_arg = False

    for i in range(first_brace + 1, len(tokens)):
        token = tokens[i]
        if token.src in START_BRACES:
            brace_stack.append(i)
        elif token.src in END_BRACES:
            brace_stack.pop()

        if len(brace_stack) == 1 and token.src == ',':
            multi_arg = True

        if not brace_stack:
            break
    else:
        raise AssertionError('Past end?')

    last_brace = i

    # This was not actually a multi-line call, despite the ast telling us that
    if tokens[first_brace].line == tokens[last_brace].line:
        return

    # determine the initial indentation
    i = first_brace
    while i >= 0 and tokens[i].name not in NEWLINES:
        i -= 1

    if i >= 0 and tokens[i + 1].name in INDENT_TOKENS:
        initial_indent = len(tokens[i + 1].src)
    else:
        initial_indent = 0

    return Fix((first_brace, last_brace), multi_arg, initial_indent)

def _find_call(call, i, tokens):
    # When we get a `call` object, the ast refers to it as this:
    #
    #     func_name(arg, arg, arg)
    #     ^ where ast points
    #
    # We care about the closing paren, in order to find it, we first walk
    # until we find an argument.  When we find an argument, we know the outer
    # paren we find is the function call paren
    #
    #     func_name(arg, arg, arg)
    #              ^ outer paren
    first_brace = None
    paren_stack = []
    for i in range(i, len(tokens)):
        token = tokens[i]
        if token.src == '(':
            paren_stack.append(i)
        # the ast lies to us about the beginning of parenthesized functions.
        # See #3. (why we make sure there's something to pop here)
        elif token.src == ')' and paren_stack:
            paren_stack.pop()

        if (token.line, token.utf8_byte_offset) in call.arg_offsets:
            first_brace = paren_stack[0]
            break
    else:
        raise AssertionError('Past end?')

    return _find_simple(first_brace, tokens)

def _find_literal(literal, i, tokens):
    # tuples are evil, we need to backtrack to find the opening paren
    if literal.backtrack:
        i -= 1
        while tokens[i].name in NON_CODING_TOKENS:
            i -= 1
        # Sometimes tuples don't even have a paren!
        # x = 1, 2, 3
        if tokens[i].src != '(':
            return

    return _find_simple(i, tokens)

def _fix_brace(fix_data, add_comma, tokens):
    first_brace, last_brace = fix_data.braces

    # Figure out if either of the braces are "hugging"
    hug_open = tokens[first_brace + 1].name not in NON_CODING_TOKENS
    hug_close = tokens[last_brace - 1].name not in NON_CODING_TOKENS
    if (
            not fix_data.multi_arg and
            hug_open and
            tokens[last_brace - 1].src in END_BRACES
    ):
        hug_open = hug_close = False

    # fix open hugging
    if hug_open:
        new_indent = fix_data.initial_indent + 4

        tokens[first_brace + 1:first_brace + 1] = [
            Token('NL', '\n'), Token(UNIMPORTANT_WS, ' ' * new_indent),
        ]
        last_brace += 2

        # Adust indentation for the rest of the things
        min_indent = None
        indents = []
        insert_indents = []
        for i in range(first_brace + 3, last_brace):
            if tokens[i - 1].name == 'NL':
                if tokens[i].name != UNIMPORTANT_WS:
                    min_indent = 0
                    insert_indents.append(i)
                else:
                    if min_indent is None:
                        min_indent = len(tokens[i].src)
                    elif len(tokens[i].src) < min_indent:
                        min_indent = len(tokens[i].src)
                    indents.append(i)

        for i in indents:
            oldlen = len(tokens[i].src)
            newlen = oldlen - min_indent + new_indent
            tokens[i] = tokens[i]._replace(src=' ' * newlen)
        for i in reversed(insert_indents):
            tokens.insert(i, Token(UNIMPORTANT_WS, ' ' * new_indent))
            last_brace += 1

    # fix close hugging
    if hug_close:
        tokens[last_brace:last_brace] = [
            Token('NL', '\n'),
            Token(UNIMPORTANT_WS, ' ' * fix_data.initial_indent),
        ]
        last_brace += 2

    # From there, we can walk backwards and decide whether a comma is needed
    i = last_brace - 1
    while tokens[i].name in NON_CODING_TOKENS:
        i -= 1

    # If we're not a hugging paren, we can insert a comma
    if add_comma and tokens[i].src != ',' and i + 1 != last_brace:
        tokens.insert(i + 1, Token('OP', ','))

    # Fix trailing brace to match leading indentation
    back_1 = tokens[last_brace - 1]
    back_2 = tokens[last_brace - 2]
    if (
            back_1.name == UNIMPORTANT_WS and
            back_2.name == 'NL' and
            len(back_1.src) != fix_data.initial_indent
    ):
        new_indent = fix_data.initial_indent * ' '
        tokens[last_brace - 1] = back_1._replace(src=new_indent)

def _changing_list(lst):
    i = 0
    while i < len(lst):
        yield i, lst[i]
        i += 1

def _fix_src(contents_text, py35_plus):
    try:
        ast_obj = ast_parse(contents_text)
    except SyntaxError:
        return contents_text

    visitor = FindNodes()
    visitor.visit(ast_obj)
    py35_plus = py35_plus or visitor.has_new_syntax

    tokens = src_to_tokens(contents_text)
    for i, token in _changing_list(tokens):
        key = Offset(token.line, token.utf8_byte_offset)
        add_comma = True
        fix_data = None

        if key in visitor.calls:
            call = visitor.calls[key]
            # Only fix stararg calls if asked to
            add_comma = not call.star_args or py35_plus
            fix_data = _find_call(call, i, tokens)
        elif key in visitor.funcs:
            func = visitor.funcs[key]
            # any amount of starargs excludes adding a comma for defs
            add_comma = not func.star_args
            # functions can be treated as calls
            fix_data = _find_call(func, i, tokens)
        # Handle parenthesized things
        elif token.src == '(':
            fix_data = _find_simple(i, tokens)
            add_comma = False

        if fix_data is not None:
            _fix_brace(fix_data, add_comma, tokens)

        # need to additionally handle literals afterwards as tuples report
        # their starting index as the first element, which may be one of the
        # above things.
        if key in visitor.literals:
            fix_data = _find_literal(visitor.literals[key], i, tokens)
            if fix_data is not None:
                _fix_brace(fix_data, True, tokens)

    return tokens_to_src(tokens)

def fix_file(filename, args):
    with open(filename, 'rb') as f:
        contents_bytes = f.read()

    try:
        contents_text_orig = contents_text = contents_bytes.decode('UTF-8')
    except UnicodeDecodeError:
        print('{} is non-utf-8 (not supported)'.format(filename))
        return 1

    contents_text = _fix_src(contents_text, args.py35_plus)

    if contents_text != contents_text_orig:
        print('Rewriting {}'.format(filename))
        with io.open(filename, 'w', encoding='UTF-8') as f:
            f.write(contents_text)
        return 1

    return 0

def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('filenames', nargs='*')
    parser.add_argument('--py35-plus', action='store_true')
    args = parser.parse_args(argv)

    ret = 0
    for filename in args.filenames:
        ret |= fix_file(filename, args)
    return ret

if __name__ == '__main__':
    exit(main())
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import subprocess

import pwndbg.commands

parser = argparse.ArgumentParser(description='Launches radare2',
                                 epilog="Example: r2 -- -S -AA")
parser.add_argument('--no-seek', action='store_true',
                    help='Do not seek to current pc')
parser.add_argument('arguments', nargs='*', type=str,
                    help='Arguments to pass to radare')

@pwndbg.commands.ArgparsedCommand(parser)
@pwndbg.commands.OnlyWithFile
def r2(arguments, no_seek=False):
    filename = pwndbg.file.get_file(pwndbg.proc.exe)

    # Build up the command line to run
    cmd = ['radare2', filename]
    if pwndbg.proc.alive:
        addr = pwndbg.regs.pc
        if pwndbg.elf.get_elf_info(filename).is_pie:
            addr -= pwndbg.elf.exe().address
        if not no_seek:
            cmd.extend(['-s', hex(addr)])
    cmd += arguments

    try:
        subprocess.call(cmd)
    except Exception:
        print("Could not run radare2. Please ensure it's installed and in $PATH.")
<commit_msg>Moved filename to the end of the command (#842)<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import subprocess

import pwndbg.commands

parser = argparse.ArgumentParser(description='Launches radare2',
                                 epilog="Example: r2 -- -S -AA")
parser.add_argument('--no-seek', action='store_true',
                    help='Do not seek to current pc')
parser.add_argument('arguments', nargs='*', type=str,
                    help='Arguments to pass to radare')

@pwndbg.commands.ArgparsedCommand(parser)
@pwndbg.commands.OnlyWithFile
def r2(arguments, no_seek=False):
    filename = pwndbg.file.get_file(pwndbg.proc.exe)

    # Build up the command line to run
    cmd = ['radare2']
    if pwndbg.proc.alive:
        addr = pwndbg.regs.pc
        if pwndbg.elf.get_elf_info(filename).is_pie:
            addr -= pwndbg.elf.exe().address
        if not no_seek:
            cmd.extend(['-s', hex(addr)])
    cmd += arguments
    cmd.extend([filename])

    try:
        subprocess.call(cmd)
    except Exception:
        print("Could not run radare2. Please ensure it's installed and in $PATH.")
<|endoftext|>


<commit_before>#
#  Copyright (C) (2013) STFC Rutherford Appleton Laboratory, UK.
#
#  Author: David Waterman.
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package.
#

# python and cctbx imports
from __future__ import division
from math import pi, sqrt
from cctbx.array_family import flex
import random

# dials imports
from dials.algorithms.refinement.target import Target
from dials.algorithms.refinement.target import ReflectionManager
from dials.algorithms.refinement.target import ObsPredMatch # implicit import

# constants
TWO_PI = 2.0 * pi

class LeastSquaresXYResidualWithRmsdCutoff(Target):
  """An implementation of the target class providing a least squares residual
  in terms of detector impact position X, Y only, terminating on achieved
  rmsd (or on intrisic convergence of the chosen minimiser)"""

  rmsd_names = ["RMSD_X", "RMSD_Y"]

  def __init__(self, experiments, reflection_predictor, ref_man,
               prediction_parameterisation,
               frac_binsize_cutoff=0.33333,
               absolute_cutoffs=None):

    Target.__init__(self, experiments, reflection_predictor, ref_man,
                    prediction_parameterisation)

    # Set up the RMSD achieved criterion. For simplicity, we take detector from
    # the first Experiment only.
    detector = experiments[0].detector
    if not absolute_cutoffs:
      pixel_sizes = [p.get_pixel_size() for p in detector]
      min_px_size_x = min(e[0] for e in pixel_sizes)
      min_px_size_y = min(e[1] for e in pixel_sizes)
      self._binsize_cutoffs = [min_px_size_x * frac_binsize_cutoff,
                               min_px_size_y * frac_binsize_cutoff]
    else:
      self._binsize_cutoffs = absolute_cutoffs[:2]

    return

  def compute_residuals_and_gradients(self):
    """return the vector of residuals plus their gradients
    and weights for non-linear least squares methods"""

    self._matches = self._reflection_manager.get_matches()

    # return residuals and weights as 1d flex.double vectors
    nelem = len(self._matches) * 2
    residuals = flex.double(nelem)
    jacobian_t = flex.double(flex.grid(
        len(self._prediction_parameterisation), nelem))
    weights = flex.double(nelem)

    for i, m in enumerate(self._matches):
      residuals[2*i] = m.x_resid
      residuals[2*i + 1] = m.y_resid
      #residuals[3*i + 2] = m.Phiresid

      # are these the right weights? Or inverse, or sqrt?
      weights[2*i] = m.weight_x_obs
      weights[2*i + 1] = m.weight_y_obs
      #weights[3*i + 2] = m.weightPhio

      # m.gradients is a nparam length list, each element of which is a
      # doublet of values, (dX/dp_n, dY/dp_n)
      dX_dp, dY_dp = zip(*m.gradients)

      # FIXME Here we paste columns into the Jacobian transpose then take
      # its transpose when done. This seems inefficient: can we just start
      # with the Jacobian and fill elements sequentially, using row-major
      # order to ensure the values are filled in the right order?

      # fill jacobian elements here.
      jacobian_t.matrix_paste_column_in_place(flex.double(dX_dp), 2*i)
      jacobian_t.matrix_paste_column_in_place(flex.double(dY_dp), 2*i+1)
      #jacobian_t.matrix_paste_column_in_place(flex.double(dPhi_dp), 3*i+2)

    # We return the Jacobian, not its transpose.
    jacobian_t.matrix_transpose_in_place()

    return(residuals, jacobian_t, weights)

  def compute_functional_and_gradients(self):
    """calculate the value of the target function and its gradients"""

    self._matches = self._reflection_manager.get_matches()
    self._nref = len(self._matches)

    # This is a hack for the case where nref=0. This should not be necessary
    # if bounds are provided for parameters to stop the algorithm exploring
    # unreasonable regions of parameter space where no predictions exist.
    # Unfortunately the L-BFGS line search does make such extreme trials.
    if self._nref == 0:
      return 1.e12, [1.] * len(self._prediction_parameterisation)

    # compute target function
    L = 0.5 * sum([m.weight_x_obs * m.x_resid2 +
                   m.weight_y_obs * m.y_resid2
                   for m in self._matches])

    # prepare list of gradients
    dL_dp = [0.] * len(self._prediction_parameterisation)

    # the gradients wrt each parameter are stored with the matches
    for m in self._matches:

      for j, (grad_X, grad_Y) in enumerate(m.gradients):
        dL_dp[j] += (m.weight_x_obs * m.x_resid * grad_X +
                     m.weight_y_obs * m.y_resid * grad_Y)

    return (L, dL_dp)

  def curvatures(self):
    """First order approximation to the diagonal of the Hessian based on the
    least squares form of the target"""

    # This is a hack for the case where nref=0. This should not be necessary
    # if bounds are provided for parameters to stop the algorithm exploring
    # unreasonable regions of parameter space where there are no predictions
    if self._nref == 0:
      return [1.] * len(self._prediction_parameterisation)

    # prepare lists of gradients and curvatures
    curv = [0.] * len(self._prediction_parameterisation)

    # for each reflection, get the approximate curvatures wrt each parameter
    for m in self._matches:

      for j, (grad_X, grad_Y) in enumerate(m.gradients):
        curv[j] += (m.weight_x_obs * grad_X**2 +
                    m.weight_y_obs * grad_Y**2)

    # Curvatures of zero will cause a crash, because their inverse is taken.
    assert all([c > 0.0 for c in curv])

    return curv

  def rmsds(self):
    """calculate unweighted RMSDs"""

    if not self._matches:
      self._matches = self._reflection_manager.get_matches()

    resid_x = sum((m.x_resid2 for m in self._matches))
    resid_y = sum((m.y_resid2 for m in self._matches))

    # cache rmsd calculation for achieved test
    n = len(self._matches)
    self._rmsds = (sqrt(resid_x / n),
                   sqrt(resid_y / n))

    return self._rmsds

  def achieved(self):
    """RMSD criterion for target achieved """
    r = self._rmsds if self._rmsds else self.rmsds()

    # reset cached rmsds to avoid getting out of step
    self._rmsds = None

    if (r[0] < self._binsize_cutoffs[0] and
        r[1] < self._binsize_cutoffs[1]):
      return True
    return False

class ReflectionManagerXY(ReflectionManager):
  """Overloads for a Reflection Manager that does not exclude
  reflections too close to the spindle, and reports only information
  about X, Y residuals"""

  def _spindle_beam_plane_normal(self):
    """There is no goniometer, so overload to return None"""

    return None

  def _id_refs_to_keep(self, obs_data):
    """For this version of the class, only reject the (0,0,0) reflections.
    We don't want to exclude reflections close to the spindle, as the spindle
    may not exist"""

    inc = [i for i, ref in enumerate(obs_data) if ref['miller_index'] != (0,0,0)]

    return inc

  def _create_working_set(self, indices, nref_per_degree,
                                         minimum_sample_size,
                                         max_num_obs):
    """Make a subset of the indices of reflections to use in refinement.

    This version ignores nref_per_degree"""

    working_indices = indices
    sample_size = len(working_indices)

    # set maximum sample size
    if max_num_obs:
      if sample_size > max_num_obs:
        sample_size = max_num_obs

    # sample the data and record the sample size
    if sample_size < len(working_indices):
      self._sample_size = sample_size
      working_indices = random.sample(working_indices,
                                      self._sample_size)
    return(working_indices)

  def print_stats_on_matches(self):
    """Print some basic statistics on the matches"""

    l = self.get_matches()

    if self._verbosity > 1:

      from scitbx.math import five_number_summary
      x_resid = [e.x_resid for e in l]
      y_resid = [e.y_resid for e in l]
      w_x = [e.weight_x_obs for e in l]
      w_y = [e.weight_y_obs for e in l]

      print "\nSummary statistics for observations matched to predictions:"
      print ("                      "
             "Min         Q1        Med         Q3        Max")
      print "(Xc-Xo)        {0:10.5g} {1:10.5g} {2:10.5g} {3:10.5g} {4:10.5g}".\
        format(*five_number_summary(x_resid))
      print "(Yc-Yo)        {0:10.5g} {1:10.5g} {2:10.5g} {3:10.5g} {4:10.5g}".\
        format(*five_number_summary(y_resid))
      print "X weights      {0:10.5g} {1:10.5g} {2:10.5g} {3:10.5g} {4:10.5g}".\
        format(*five_number_summary(w_x))
      print "Y weights      {0:10.5g} {1:10.5g} {2:10.5g} {3:10.5g} {4:10.5g}".\
        format(*five_number_summary(w_y))
      print

      if len(l) >= 20 and self._verbosity > 2:

        sl = self._sort_obs_by_residual(l)
        print "Reflections with the worst 20 positional residuals:"
        print "H, K, L, x_resid, y_resid, weight_x_obs, weight_y_obs"
        fmt = "(%3d, %3d, %3d) %5.3f %5.3f %5.3f %5.3f"
        for i in xrange(20):
          e = sl[i]
          msg = fmt % tuple(e.miller_index + (e.x_resid,
                           e.y_resid,
                           e.weight_x_obs,
                           e.weight_y_obs))
          print msg
        print

    return

  def reject_outliers(self):
    """Unset the use flag on matches with extreme (outlying) residuals.

    Outlier detection finds values more than _iqr_multiplier times the
    interquartile range from the quartiles. When x=1.5, this is Tukey's rule.

    Return boolean whether rejection was performed or not"""

    # return early if outlier rejection is disabled
    if self._iqr_multiplier is None: return False

    from scitbx.math import five_number_summary
    matches = [obs for obs in self._obs_pred_pairs if obs.is_matched]

    x_resid = [e.x_resid for e in matches]
    y_resid = [e.y_resid for e in matches]

    min_x, q1_x, med_x, q3_x, max_x = five_number_summary(x_resid)
    min_y, q1_y, med_y, q3_y, max_y = five_number_summary(y_resid)

    iqr_x = q3_x - q1_x
    iqr_y = q3_y - q1_y

    cut_x = self._iqr_multiplier * iqr_x
    cut_y = self._iqr_multiplier * iqr_y

    for m in matches:
      if m.x_resid > q3_x + cut_x: m.is_matched = False
      if m.x_resid < q1_x - cut_x: m.is_matched = False
      if m.y_resid > q3_y + cut_y: m.is_matched = False
      if m.y_resid < q1_y - cut_y: m.is_matched = False

    nreject = [m.is_matched for m in matches].count(False)

    if nreject == 0: return False

    if self._verbosity > 1:
      print "%d reflections have been rejected as outliers" % nreject

    return True
<commit_msg>change names in line with other versions of the class<commit_after>#
#  Copyright (C) (2013) STFC Rutherford Appleton Laboratory, UK.
#
#  Author: David Waterman.
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package.
#

# python and cctbx imports
from __future__ import division
from math import pi, sqrt
from cctbx.array_family import flex
import random

# dials imports
from dials.algorithms.refinement.target import Target
from dials.algorithms.refinement.target import ReflectionManager
from dials.algorithms.refinement.target import ObsPredMatch # implicit import

# constants
TWO_PI = 2.0 * pi

class LeastSquaresXYResidualWithRmsdCutoff(Target):
  """An implementation of the target class providing a least squares residual
  in terms of detector impact position X, Y only, terminating on achieved
  rmsd (or on intrisic convergence of the chosen minimiser)"""

  rmsd_names = ["RMSD_X", "RMSD_Y"]

  def __init__(self, experiments, reflection_predictor, ref_man,
               prediction_parameterisation,
               frac_binsize_cutoff=0.33333,
               absolute_cutoffs=None):

    Target.__init__(self, experiments, reflection_predictor, ref_man,
                    prediction_parameterisation)

    # Set up the RMSD achieved criterion. For simplicity, we take detector from
    # the first Experiment only.
    detector = experiments[0].detector
    if not absolute_cutoffs:
      pixel_sizes = [p.get_pixel_size() for p in detector]
      min_px_size_x = min(e[0] for e in pixel_sizes)
      min_px_size_y = min(e[1] for e in pixel_sizes)
      self._binsize_cutoffs = [min_px_size_x * frac_binsize_cutoff,
                               min_px_size_y * frac_binsize_cutoff]
    else:
      self._binsize_cutoffs = absolute_cutoffs[:2]

    return

  def compute_residuals_and_gradients(self):
    """return the vector of residuals plus their gradients
    and weights for non-linear least squares methods"""

    self._matches = self._reflection_manager.get_matches()

    # return residuals and weights as 1d flex.double vectors
    nelem = len(self._matches) * 2
    residuals = flex.double(nelem)
    jacobian_t = flex.double(flex.grid(
        len(self._prediction_parameterisation), nelem))
    weights = flex.double(nelem)

    for i, m in enumerate(self._matches):
      residuals[2*i] = m.x_resid
      residuals[2*i + 1] = m.y_resid
      #residuals[3*i + 2] = m.Phiresid

      # are these the right weights? Or inverse, or sqrt?
      weights[2*i] = m.weight_x_obs
      weights[2*i + 1] = m.weight_y_obs
      #weights[3*i + 2] = m.weightPhio

      # m.gradients is a nparam length list, each element of which is a
      # doublet of values, (dX/dp_n, dY/dp_n)
      dX_dp, dY_dp = zip(*m.gradients)

      # FIXME Here we paste columns into the Jacobian transpose then take
      # its transpose when done. This seems inefficient: can we just start
      # with the Jacobian and fill elements sequentially, using row-major
      # order to ensure the values are filled in the right order?

      # fill jacobian elements here.
      jacobian_t.matrix_paste_column_in_place(flex.double(dX_dp), 2*i)
      jacobian_t.matrix_paste_column_in_place(flex.double(dY_dp), 2*i+1)
      #jacobian_t.matrix_paste_column_in_place(flex.double(dPhi_dp), 3*i+2)

    # We return the Jacobian, not its transpose.
    jacobian_t.matrix_transpose_in_place()

    return(residuals, jacobian_t, weights)

  def compute_functional_and_gradients(self):
    """calculate the value of the target function and its gradients"""

    self._matches = self._reflection_manager.get_matches()
    self._nref = len(self._matches)

    # This is a hack for the case where nref=0. This should not be necessary
    # if bounds are provided for parameters to stop the algorithm exploring
    # unreasonable regions of parameter space where no predictions exist.
    # Unfortunately the L-BFGS line search does make such extreme trials.
    if self._nref == 0:
      return 1.e12, [1.] * len(self._prediction_parameterisation)

    # compute target function
    L = 0.5 * sum([m.weight_x_obs * m.x_resid2 +
                   m.weight_y_obs * m.y_resid2
                   for m in self._matches])

    # prepare list of gradients
    dL_dp = [0.] * len(self._prediction_parameterisation)

    # the gradients wrt each parameter are stored with the matches
    for m in self._matches:

      for j, (grad_X, grad_Y) in enumerate(m.gradients):
        dL_dp[j] += (m.weight_x_obs * m.x_resid * grad_X +
                     m.weight_y_obs * m.y_resid * grad_Y)

    return (L, dL_dp)

  def curvatures(self):
    """First order approximation to the diagonal of the Hessian based on the
    least squares form of the target"""

    # This is a hack for the case where nref=0. This should not be necessary
    # if bounds are provided for parameters to stop the algorithm exploring
    # unreasonable regions of parameter space where there are no predictions
    if self._nref == 0:
      return [1.] * len(self._prediction_parameterisation)

    # prepare lists of gradients and curvatures
    curv = [0.] * len(self._prediction_parameterisation)

    # for each reflection, get the approximate curvatures wrt each parameter
    for m in self._matches:

      for j, (grad_X, grad_Y) in enumerate(m.gradients):
        curv[j] += (m.weight_x_obs * grad_X**2 +
                    m.weight_y_obs * grad_Y**2)

    # Curvatures of zero will cause a crash, because their inverse is taken.
    assert all([c > 0.0 for c in curv])

    return curv

  def rmsds(self):
    """calculate unweighted RMSDs"""

    if not self._matches:
      self._matches = self._reflection_manager.get_matches()

    resid_x = sum((m.x_resid2 for m in self._matches))
    resid_y = sum((m.y_resid2 for m in self._matches))

    # cache rmsd calculation for achieved test
    n = len(self._matches)
    self._rmsds = (sqrt(resid_x / n),
                   sqrt(resid_y / n))

    return self._rmsds

  def achieved(self):
    """RMSD criterion for target achieved """
    r = self._rmsds if self._rmsds else self.rmsds()

    # reset cached rmsds to avoid getting out of step
    self._rmsds = None

    if (r[0] < self._binsize_cutoffs[0] and
        r[1] < self._binsize_cutoffs[1]):
      return True
    return False

class ReflectionManagerXY(ReflectionManager):
  """Overloads for a Reflection Manager that does not exclude
  reflections too close to the spindle, and reports only information
  about X, Y residuals"""

  def _spindle_beam_plane_normal(self):
    """There is no goniometer, so overload to return None"""

    return None

  def _id_refs_to_keep(self, obs_data):
    """For this version of the class, only reject the (0,0,0) reflections.
    We don't want to exclude reflections close to the spindle, as the spindle
    may not exist"""

    inc = [i for i, ref in enumerate(obs_data) if ref['miller_index'] != (0,0,0)]

    return inc

  def _create_working_set(self, indices, nref_per_degree,
                                         min_sample_size,
                                         max_sample_size):
    """Make a subset of the indices of reflections to use in refinement.

    This version ignores nref_per_degree"""

    working_indices = indices
    sample_size = len(working_indices)

    # set maximum sample size
    if max_sample_size:
      if sample_size > max_sample_size:
        sample_size = max_sample_size

    # sample the data and record the sample size
    if sample_size < len(working_indices):
      self._sample_size = sample_size
      working_indices = random.sample(working_indices,
                                      self._sample_size)
    return(working_indices)

  def print_stats_on_matches(self):
    """Print some basic statistics on the matches"""

    l = self.get_matches()

    if self._verbosity > 1:

      from scitbx.math import five_number_summary
      x_resid = [e.x_resid for e in l]
      y_resid = [e.y_resid for e in l]
      w_x = [e.weight_x_obs for e in l]
      w_y = [e.weight_y_obs for e in l]

      print "\nSummary statistics for observations matched to predictions:"
      print ("                      "
             "Min         Q1        Med         Q3        Max")
      print "(Xc-Xo)        {0:10.5g} {1:10.5g} {2:10.5g} {3:10.5g} {4:10.5g}".\
        format(*five_number_summary(x_resid))
      print "(Yc-Yo)        {0:10.5g} {1:10.5g} {2:10.5g} {3:10.5g} {4:10.5g}".\
        format(*five_number_summary(y_resid))
      print "X weights      {0:10.5g} {1:10.5g} {2:10.5g} {3:10.5g} {4:10.5g}".\
        format(*five_number_summary(w_x))
      print "Y weights      {0:10.5g} {1:10.5g} {2:10.5g} {3:10.5g} {4:10.5g}".\
        format(*five_number_summary(w_y))
      print

      if len(l) >= 20 and self._verbosity > 2:

        sl = self._sort_obs_by_residual(l)
        print "Reflections with the worst 20 positional residuals:"
        print "H, K, L, x_resid, y_resid, weight_x_obs, weight_y_obs"
        fmt = "(%3d, %3d, %3d) %5.3f %5.3f %5.3f %5.3f"
        for i in xrange(20):
          e = sl[i]
          msg = fmt % tuple(e.miller_index + (e.x_resid,
                           e.y_resid,
                           e.weight_x_obs,
                           e.weight_y_obs))
          print msg
        print

    return

  def reject_outliers(self):
    """Unset the use flag on matches with extreme (outlying) residuals.

    Outlier detection finds values more than _iqr_multiplier times the
    interquartile range from the quartiles. When x=1.5, this is Tukey's rule.

    Return boolean whether rejection was performed or not"""

    # return early if outlier rejection is disabled
    if self._iqr_multiplier is None: return False

    from scitbx.math import five_number_summary
    matches = [obs for obs in self._obs_pred_pairs if obs.is_matched]

    x_resid = [e.x_resid for e in matches]
    y_resid = [e.y_resid for e in matches]

    min_x, q1_x, med_x, q3_x, max_x = five_number_summary(x_resid)
    min_y, q1_y, med_y, q3_y, max_y = five_number_summary(y_resid)

    iqr_x = q3_x - q1_x
    iqr_y = q3_y - q1_y

    cut_x = self._iqr_multiplier * iqr_x
    cut_y = self._iqr_multiplier * iqr_y

    for m in matches:
      if m.x_resid > q3_x + cut_x: m.is_matched = False
      if m.x_resid < q1_x - cut_x: m.is_matched = False
      if m.y_resid > q3_y + cut_y: m.is_matched = False
      if m.y_resid < q1_y - cut_y: m.is_matched = False

    nreject = [m.is_matched for m in matches].count(False)

    if nreject == 0: return False

    if self._verbosity > 1:
      print "%d reflections have been rejected as outliers" % nreject

    return True
<|endoftext|>