<commit_before># This file is part of Maker Keeper Framework.
#
# Copyright (C) 2017 reverendus
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import json
from typing import Optional

import pkg_resources
from web3 import Web3, EthereumTesterProvider

from pymaker import Address
from pymaker.approval import directly
from pymaker.auth import DSGuard
from pymaker.etherdelta import EtherDelta
from pymaker.feed import DSValue
from pymaker.numeric import Wad
from pymaker.oasis import MatchingMarket
from pymaker.sai import Tub, Tap, Top
from pymaker.token import DSToken
from pymaker.vault import DSVault

def deploy_contract(web3: Web3, contract_name: str, args: Optional[list]=None) -> Address:
    """Deploys a new contract.

    Args:
        web3: An instance of `Web` from `web3.py`.
        contract_name: Name of the contract, used to find the `abi` and `bin` files.
        args: Optional list of contract constructor parameters.

    Returns:
        Ethereum address of the newly deployed contract, as a :py:class:`pymaker.Address` instance.
    """
    assert(isinstance(web3, Web3))
    assert(isinstance(contract_name, str))
    assert(isinstance(args, list) or (args is None))

    abi = json.loads(pkg_resources.resource_string('pymaker.deployment', f'abi/{contract_name}.abi'))
    bytecode = pkg_resources.resource_string('pymaker.deployment', f'abi/{contract_name}.bin')
    tx_hash = web3.eth.contract(abi=abi, bytecode=bytecode).deploy(args=args)
    receipt = web3.eth.getTransactionReceipt(tx_hash)
    return Address(receipt['contractAddress'])

class Deployment:
    """Represents a test deployment of the entire Maker smart contract ecosystem.

    Creating an instance of this class creates a testrpc web3 provider with the entire set
    of Maker smart contracts deployed to it. It is used in unit tests of PyMaker, and also in
    unit tests for individual keepers.
    """
    def __init__(self):
        web3 = Web3(EthereumTesterProvider())
        web3.eth.defaultAccount = web3.eth.accounts[0]
        our_address = Address(web3.eth.defaultAccount)
        sai = DSToken.deploy(web3, 'SAI')
        sin = DSToken.deploy(web3, 'SIN')
        gem = DSToken.deploy(web3, 'ETH')
        pip = DSValue.deploy(web3)
        skr = DSToken.deploy(web3, 'SKR')
        pot = DSVault.deploy(web3)
        pit = DSVault.deploy(web3)
        tip = deploy_contract(web3, 'Tip')
        dad = DSGuard.deploy(web3)
        jug = deploy_contract(web3, 'SaiJug', [sai.address.address, sin.address.address])
        jar = deploy_contract(web3, 'SaiJar', [skr.address.address, gem.address.address, pip.address.address])

        tub = Tub.deploy(web3, jar, jug, pot.address, pit.address, tip)
        tap = Tap.deploy(web3, tub.address, pit.address)
        top = Top.deploy(web3, tub.address, tap.address)

        otc = MatchingMarket.deploy(web3, 2600000000)
        etherdelta = EtherDelta.deploy(web3,
                                       admin=Address('0x1111100000999998888877777666665555544444'),
                                       fee_account=Address('0x8888877777666665555544444111110000099999'),
                                       account_levels_addr=Address('0x0000000000000000000000000000000000000000'),
                                       fee_make=Wad.from_number(0.01),
                                       fee_take=Wad.from_number(0.02),
                                       fee_rebate=Wad.from_number(0.03))

        # set permissions
        dad.permit(DSGuard.ANY, DSGuard.ANY, DSGuard.ANY).transact()
        tub.set_authority(dad.address)
        for auth in [sai, sin, skr, pot, pit, tap, top]:
            auth.set_authority(dad.address).transact()

        # whitelist pairs
        otc.add_token_pair_whitelist(sai.address, gem.address).transact()

        # approve, mint some GEMs
        tub.approve(directly())
        gem.mint(Wad.from_number(1000000)).transact()

        web3.providers[0].rpc_methods.evm_snapshot()

        self.web3 = web3
        self.our_address = our_address
        self.gem = gem
        self.sai = sai
        self.sin = sin
        self.skr = skr
        self.tub = tub
        self.tap = tap
        self.top = top
        self.otc = otc
        self.etherdelta = etherdelta

    def reset(self):
        """Rollbacks all changes made since the initial deployment."""
        self.web3.providers[0].rpc_methods.evm_revert()
        self.web3.providers[0].rpc_methods.evm_snapshot()
        self.otc._none_orders = set()
<commit_msg>`Deployment` class reflecting SAI2.0/DAI contract changes<commit_after># This file is part of Maker Keeper Framework.
#
# Copyright (C) 2017 reverendus
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import json
from typing import Optional

import pkg_resources
from web3 import Web3, EthereumTesterProvider

from pymaker import Address
from pymaker.approval import directly
from pymaker.auth import DSGuard
from pymaker.etherdelta import EtherDelta
from pymaker.feed import DSValue
from pymaker.numeric import Wad
from pymaker.oasis import MatchingMarket
from pymaker.sai import Tub, Tap, Top
from pymaker.token import DSToken
from pymaker.vault import DSVault

def deploy_contract(web3: Web3, contract_name: str, args: Optional[list]=None) -> Address:
    """Deploys a new contract.

    Args:
        web3: An instance of `Web` from `web3.py`.
        contract_name: Name of the contract, used to find the `abi` and `bin` files.
        args: Optional list of contract constructor parameters.

    Returns:
        Ethereum address of the newly deployed contract, as a :py:class:`pymaker.Address` instance.
    """
    assert(isinstance(web3, Web3))
    assert(isinstance(contract_name, str))
    assert(isinstance(args, list) or (args is None))

    abi = json.loads(pkg_resources.resource_string('pymaker.deployment', f'abi/{contract_name}.abi'))
    bytecode = pkg_resources.resource_string('pymaker.deployment', f'abi/{contract_name}.bin')
    tx_hash = web3.eth.contract(abi=abi, bytecode=bytecode).deploy(args=args)
    receipt = web3.eth.getTransactionReceipt(tx_hash)
    return Address(receipt['contractAddress'])

class Deployment:
    """Represents a test deployment of the entire Maker smart contract ecosystem.

    Creating an instance of this class creates a testrpc web3 provider with the entire set
    of Maker smart contracts deployed to it. It is used in unit tests of PyMaker, and also in
    unit tests for individual keepers.
    """
    def __init__(self):
        web3 = Web3(EthereumTesterProvider())
        web3.eth.defaultAccount = web3.eth.accounts[0]
        our_address = Address(web3.eth.defaultAccount)
        sai = DSToken.deploy(web3, 'SAI')
        sin = DSToken.deploy(web3, 'SIN')
        skr = DSToken.deploy(web3, 'SKR')
        gem = DSToken.deploy(web3, 'ETH')
        gov = DSToken.deploy(web3, 'MKR')
        pip = DSValue.deploy(web3)
        pep = DSValue.deploy(web3)
        pit = DSVault.deploy(web3)

        vox = deploy_contract(web3, 'SaiVox')
        tub = Tub.deploy(web3, sai=sai.address, sin=sin.address, skr=skr.address, gem=gem.address, gov=gov.address,
                         pip=pip.address, pep=pep.address, vox=vox, pit=pit.address)
        tap = Tap.deploy(web3, tub.address)
        top = Top.deploy(web3, tub.address, tap.address)

        otc = MatchingMarket.deploy(web3, 2600000000)
        etherdelta = EtherDelta.deploy(web3,
                                       admin=Address('0x1111100000999998888877777666665555544444'),
                                       fee_account=Address('0x8888877777666665555544444111110000099999'),
                                       account_levels_addr=Address('0x0000000000000000000000000000000000000000'),
                                       fee_make=Wad.from_number(0.01),
                                       fee_take=Wad.from_number(0.02),
                                       fee_rebate=Wad.from_number(0.03))

        # set permissions
        dad = DSGuard.deploy(web3)
        dad.permit(DSGuard.ANY, DSGuard.ANY, DSGuard.ANY).transact()
        tub.set_authority(dad.address)
        for auth in [sai, sin, skr, gem, gov, pit, tap, top]:
            auth.set_authority(dad.address).transact()

        # whitelist pairs
        otc.add_token_pair_whitelist(sai.address, gem.address).transact()

        # approve, mint some GEMs
        tub.approve(directly())
        gem.mint(Wad.from_number(1000000)).transact()

        web3.providers[0].rpc_methods.evm_snapshot()

        self.web3 = web3
        self.our_address = our_address
        self.sai = sai
        self.sin = sin
        self.skr = skr
        self.gem = gem
        self.gov = gov
        self.tub = tub
        self.tap = tap
        self.top = top
        self.otc = otc
        self.etherdelta = etherdelta

    def reset(self):
        """Rollbacks all changes made since the initial deployment."""
        self.web3.providers[0].rpc_methods.evm_revert()
        self.web3.providers[0].rpc_methods.evm_snapshot()
        self.otc._none_orders = set()
<|endoftext|>


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

import re

from .board import *
from .pieces_act import PIECE_TO_ACT, TURN_PIECE_REVERSED, KANJI_TO_PIECE, TEAITXT_TO_TEAI

SYMBOL_TO_CODE = {'▲': TEBAN_CODE[0], '△': TEBAN_CODE[1]}
KANJI_TO_INT = dict(zip(tuple('一二三四五六七八九'), range(1, 10)))
ZEN_TO_INT = dict(zip(tuple('１２３４５６７８９'), range(1, 10)))
REGEX_MOVE = re.compile('([▲△](?:同\u3000)?[^▲△\s]+)')
OGOMA = ['HI', 'KA', 'RY', 'UM']

TO_INT = dict(zip([str(r) for r in range(1, 10)], range(1, 10)))
TO_INT.update(KANJI_TO_INT)
TO_INT.update(ZEN_TO_INT)

class Ki2converter:

    '''Converter from KI2 format to CSA format.
    '''

    def __init__(self):
        self.ki2_txt = ''
        self.moves_ki2 = []
        self.comments = []
        self.infos = []

    def from_path(self, ki2path):

        try:
            with open(ki2path, 'r', encoding='shiftjis') as f:
                lines = f.readlines()
        except UnicodeDecodeError:
            with open(ki2path, 'r', encoding='CP932') as f:
                lines = f.readlines()

        self.from_lines(lines)
        self.ki2_txt = ''.join(lines)

    def from_txt(self, ki2_txt):
        '''

        ki2_txt : str
            str object of KI2 format
        '''
        self.ki2_txt = ki2_txt
        self.from_lines(ki2_txt.split('\n'))

    def from_lines(self, lines):
        moves = []
        comments = []
        infos = []

        for line in lines:
            # Comment line
            if line.startswith('*'):
                comments.append(line[1:])

            # Move line
            elif line.startswith('▲') or line.startswith('△'):
                line = re.sub('  ', '\u3000', line)
                tesu_line = len(re.findall('[▲△]', line))

                if tesu_line >= 2:
                    moves_line = re.findall(REGEX_MOVE, line)
                    moves.extend(moves_line)

                elif len(line) > 3:
                    moves.append(line)

            # Header information
            else:
                infos.append(line)

        self.moves_ki2 = [re.sub('(\u3000|\n)', '', move) for move in moves]
        while '' in self.moves_ki2:
            self.moves_ki2.remove('')

        self.infos = ''.join(infos)
        self.comments = ''.join(comments)

    def extract_header_infos(self):
        header_infos = {'senteban': '+'}
        queries = [
            ['sente', '^\s*[先下]手：(.+)'],
            ['gote', '^\s*[後上]手：(.+)'],
            ['start_time', '^\s*開始日時：(.+)'],
            ['end_time', '^\s*終了日時：(.+)'],
            ['site', '^\s*場所：(.+)'],
            ['event', '^\s*棋戦：(.+)'],
            ['opening', '^\s*戦型：(.+)'],
            ['time', '^\s*持ち時間：(.+)']
        ]
        for query in queries:
            match = re.findall(query[1], self.ki2_txt, re.M)

            if match:
                header_infos[query[0]] = match[0]

        # Detect teai
        for t_teaitxt, t_teai in TEAITXT_TO_TEAI.items():
            if t_teaitxt in self.ki2_txt:
                teai = t_teai
                break
        else:
            if not '手合割' in self.ki2_txt:
                teai = 'hirate'
            else:
                raise RuntimeError('invalid teai')

        header_infos['teai'] = teai
        if teai == 'sonota':
            # とりあえず
            return -1

        elif teai != 'hirate':
            header_infos['csa_komaochi'] = 'PI' + ''.join(KOMAOCHI_OPTIONS[teai])
            header_infos['senteban'] = '-'

        if not ('sente' in header_infos and 'gote' in header_infos):
            return -1

        return header_infos

    def to_csa(self):

        self.board = Board()
        header_infos = self.extract_header_infos()

        if header_infos == -1:
            return None

        self.board.set_initial_state(teai=header_infos['teai'])

        # とりあえず
        if re.search('手合割：.{,2}落ち', self.ki2_txt):
            komaochi_txt = '\n' + header_infos['csa_komaochi']
        else:
            komaochi_txt = ' 初期配置'

        csa_kifu = [
            '\'バージョン',
            'V2.2',
            '\'対局者名',
            'N+', header_infos['sente'],
            'N-', header_infos['gote'],
            '\'棋譜情報',
            '\'棋戦名',
            '$EVENT:' + header_infos.get('event', ''),
            '\'対局場所',
            '$SITE:' + header_infos.get('site', ''),
            '\'開始日時',
            '$START_TIME:' + header_infos.get('start_time', ''),
            '\'持ち時間:' + header_infos.get('time', ''),
            '\'$TIME_LIMIT:',
            '\'開始局面' + komaochi_txt,
            '\'先手番', header_infos['senteban'],
            '\'指し手と消費時間'
        ]

        # Move pieces
        for move_ki2 in self.moves_ki2:
            move_csa = self.move_ki2_to_csa(move_ki2)

            # If move_ki2_to_csa cannot read `move_ki2`
            if not move_csa:
                return None

            self.board.move(move_csa)

            csa_kifu.append(move_csa)
            csa_kifu.append('T0')

        if self.ki2_txt.endswith('勝ち'):
            csa_kifu.append('%TORYO')
        elif self.ki2_txt.endswith('千日手'):
            csa_kifu.append('%SENNICHITE')
        else:
            return None

        return '\n'.join(csa_kifu)

    def move_ki2_to_csa(self, move_ki2):
        '''Convert codes of KI2 format to that of CSA format
        '''
        replace_list = [
            ['成桂', '圭'],
            ['成香', '杏'],
            ['成銀', '全'],
            ['龍',   '竜']
        ]
        for l in replace_list:
            move_ki2 = move_ki2.replace(l[0], l[1])

        regex = re.compile('([▲△])(同|[\d一二三四五六七八九]{2})(.)(.+)?')

        reg = re.search(regex, move_ki2)

        if reg:
            matched = reg.groups()
        else:
            return None

        code = SYMBOL_TO_CODE[matched[0]]

        if matched[1] == '同':
            i = int(self.board.last_move_txt[3]) - 1
            j = int(self.board.last_move_txt[4]) - 1
        else:
            i = TO_INT[matched[1][0]] - 1
            j = TO_INT[matched[1][1]] - 1

        piece = KANJI_TO_PIECE[matched[2]]

        teban = int(code == '+')
        direction = 2 * teban - 1

        pieces_index = self.board.get_piece_indexes(piece)[1 - teban]
        prev_pos_candidates = []

        # Search candidate pieces that can move to (i, j)
        for pi, pj in pieces_index:
            koma_act = PIECE_TO_ACT[piece]
            for act in koma_act:
                for move in act:

                    next_i = pi + move[0] * direction
                    next_j = pj + move[1] * direction

                    # If next_i or next_j is outside of the board
                    if (next_i < 0 or 9 <= next_i or
                            next_j < 0 or 9 <= next_j):
                        break

                    # If conflict with the piece of move_ki2
                    if i == next_i and j == next_j:
                        prev_pos_candidates.append([pi, pj])

                    # If conflict with other pieces
                    if self.board[next_i][next_j] != EMPTY_STR:
                        break

        # If promote
        if matched[3] and re.findall('(?<!不)成', matched[3]):
            piece = TURN_PIECE_REVERSED[piece]

        if matched[3] and '打' in matched[3]:
            prev_pos_candidates = [[-1, -1]]

        # Put piece from mochigoma
        if not prev_pos_candidates:
            prev_pos_candidates = [[-1, -1]]

        n_candidates = len(prev_pos_candidates)
        if n_candidates == 1:
            prev_pos = prev_pos_candidates[0]

        elif n_candidates >= 2:
            if matched[3]:
                if '右' in matched[3]:
                    if piece in OGOMA:
                        prev_pos_candidates = filter(
                            lambda x: (i - x[0]) * direction >= 0,
                            prev_pos_candidates)
                    else:
                        prev_pos_candidates = filter(
                            lambda x: (i - x[0]) * direction > 0,
                            prev_pos_candidates)
                elif '左' in matched[3]:
                    if piece in OGOMA:
                        prev_pos_candidates = filter(
                            lambda x: (i - x[0]) * direction <= 0,
                            prev_pos_candidates)
                    else:
                        prev_pos_candidates = filter(
                            lambda x: (i - x[0]) * direction < 0,
                            prev_pos_candidates)

                if '上' in matched[3] or '行' in matched[3]:
                    prev_pos_candidates = filter(
                        lambda x: (x[1] - j) * direction > 0,
                        prev_pos_candidates
                    )
                elif '直' in matched[3]:
                    prev_pos_candidates = filter(
                        lambda x: x[0] == i and (x[1] - j) * direction > 0,
                        prev_pos_candidates
                    )
                elif '寄' in matched[3]:
                    prev_pos_candidates = filter(
                        lambda x: x[1] == j,
                        prev_pos_candidates
                    )
                elif '引' in matched[3]:
                    prev_pos_candidates = filter(
                        lambda x: (x[1] - j) * direction < 0,
                        prev_pos_candidates
                    )

                # cf. http://www.shogi.or.jp/faq/kihuhyouki.html
                # "竜が2枚の場合はやはり動作を優先します。
                #  ただし、「直」は使わずに「左」「右」で記入します。"
                if (piece in OGOMA and
                    ('右' in matched[3] or '左' in matched[3])
                ):
                    if '右' in matched[3]:
                        lr = -1
                    elif '左' in matched[3]:
                        lr = 1

                    prev_pos_candidates = [
                        max(prev_pos_candidates,
                            key=lambda x: lr * direction * x[0])
                    ]

                prev_pos_candidates = list(prev_pos_candidates)
                if (not prev_pos_candidates or
                        len(prev_pos_candidates) >= 2):
                    raise RuntimeError('Parse Error')

                prev_pos = prev_pos_candidates[0]
        else:
            raise RuntimeError('Cannot find prev position.')

        move_csa = ''.join([
            code,
            str(prev_pos[0] + 1),
            str(prev_pos[1] + 1),
            str(i + 1),
            str(j + 1),
            piece
        ])

        return move_csa
<commit_msg>Fix about prev_pos<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import re

from .board import *
from .pieces_act import PIECE_TO_ACT, TURN_PIECE_REVERSED, KANJI_TO_PIECE, TEAITXT_TO_TEAI

SYMBOL_TO_CODE = {'▲': TEBAN_CODE[0], '△': TEBAN_CODE[1]}
KANJI_TO_INT = dict(zip(tuple('一二三四五六七八九'), range(1, 10)))
ZEN_TO_INT = dict(zip(tuple('１２３４５６７８９'), range(1, 10)))
REGEX_MOVE = re.compile('([▲△](?:同\u3000)?[^▲△\s]+)')
OGOMA = ['HI', 'KA', 'RY', 'UM']

TO_INT = dict(zip([str(r) for r in range(1, 10)], range(1, 10)))
TO_INT.update(KANJI_TO_INT)
TO_INT.update(ZEN_TO_INT)

class Ki2converter:

    '''Converter from KI2 format to CSA format.
    '''

    def __init__(self):
        self.ki2_txt = ''
        self.moves_ki2 = []
        self.comments = []
        self.infos = []

    def from_path(self, ki2path):

        try:
            with open(ki2path, 'r', encoding='shiftjis') as f:
                lines = f.readlines()
        except UnicodeDecodeError:
            with open(ki2path, 'r', encoding='CP932') as f:
                lines = f.readlines()

        self.from_lines(lines)
        self.ki2_txt = ''.join(lines)

    def from_txt(self, ki2_txt):
        '''

        ki2_txt : str
            str object of KI2 format
        '''
        self.ki2_txt = ki2_txt
        self.from_lines(ki2_txt.split('\n'))

    def from_lines(self, lines):
        moves = []
        comments = []
        infos = []

        for line in lines:
            # Comment line
            if line.startswith('*'):
                comments.append(line[1:])

            # Move line
            elif line.startswith('▲') or line.startswith('△'):
                line = re.sub('  ', '\u3000', line)
                tesu_line = len(re.findall('[▲△]', line))

                if tesu_line >= 2:
                    moves_line = re.findall(REGEX_MOVE, line)
                    moves.extend(moves_line)

                elif len(line) > 3:
                    moves.append(line)

            # Header information
            else:
                infos.append(line)

        self.moves_ki2 = [re.sub('(\u3000|\n)', '', move) for move in moves]
        while '' in self.moves_ki2:
            self.moves_ki2.remove('')

        self.infos = ''.join(infos)
        self.comments = ''.join(comments)

    def extract_header_infos(self):
        header_infos = {'senteban': '+'}
        queries = [
            ['sente', '^\s*[先下]手：(.+)'],
            ['gote', '^\s*[後上]手：(.+)'],
            ['start_time', '^\s*開始日時：(.+)'],
            ['end_time', '^\s*終了日時：(.+)'],
            ['site', '^\s*場所：(.+)'],
            ['event', '^\s*棋戦：(.+)'],
            ['opening', '^\s*戦型：(.+)'],
            ['time', '^\s*持ち時間：(.+)']
        ]
        for query in queries:
            match = re.findall(query[1], self.ki2_txt, re.M)

            if match:
                header_infos[query[0]] = match[0]

        # Detect teai
        for t_teaitxt, t_teai in TEAITXT_TO_TEAI.items():
            if t_teaitxt in self.ki2_txt:
                teai = t_teai
                break
        else:
            if not '手合割' in self.ki2_txt:
                teai = 'hirate'
            else:
                raise RuntimeError('invalid teai')

        header_infos['teai'] = teai
        if teai == 'sonota':
            # とりあえず
            return -1

        elif teai != 'hirate':
            header_infos['csa_komaochi'] = 'PI' + ''.join(KOMAOCHI_OPTIONS[teai])
            header_infos['senteban'] = '-'

        if not ('sente' in header_infos and 'gote' in header_infos):
            return -1

        return header_infos

    def to_csa(self):

        self.board = Board()
        header_infos = self.extract_header_infos()

        if header_infos == -1:
            return None

        self.board.set_initial_state(teai=header_infos['teai'])

        # とりあえず
        if re.search('手合割：.{,2}落ち', self.ki2_txt):
            komaochi_txt = '\n' + header_infos['csa_komaochi']
        else:
            komaochi_txt = ' 初期配置'

        csa_kifu = [
            '\'バージョン',
            'V2.2',
            '\'対局者名',
            'N+', header_infos['sente'],
            'N-', header_infos['gote'],
            '\'棋譜情報',
            '\'棋戦名',
            '$EVENT:' + header_infos.get('event', ''),
            '\'対局場所',
            '$SITE:' + header_infos.get('site', ''),
            '\'開始日時',
            '$START_TIME:' + header_infos.get('start_time', ''),
            '\'持ち時間:' + header_infos.get('time', ''),
            '\'$TIME_LIMIT:',
            '\'開始局面' + komaochi_txt,
            '\'先手番', header_infos['senteban'],
            '\'指し手と消費時間'
        ]

        # Move pieces
        for move_ki2 in self.moves_ki2:
            move_csa = self.move_ki2_to_csa(move_ki2)

            # If move_ki2_to_csa cannot read `move_ki2`
            if not move_csa:
                return None

            self.board.move(move_csa)

            csa_kifu.append(move_csa)
            csa_kifu.append('T0')

        if self.ki2_txt.endswith('勝ち'):
            csa_kifu.append('%TORYO')
        elif self.ki2_txt.endswith('千日手'):
            csa_kifu.append('%SENNICHITE')
        else:
            return None

        return '\n'.join(csa_kifu)

    def move_ki2_to_csa(self, move_ki2):
        '''Convert codes of KI2 format to that of CSA format
        '''
        prev_pos = None
        replace_list = [
            ['成桂', '圭'],
            ['成香', '杏'],
            ['成銀', '全'],
            ['龍',   '竜']
        ]
        for l in replace_list:
            move_ki2 = move_ki2.replace(l[0], l[1])

        regex = re.compile('([▲△])(同|[\d一二三四五六七八九]{2})(.)(.+)?')

        reg = re.search(regex, move_ki2)

        if reg:
            matched = reg.groups()
        else:
            return None

        code = SYMBOL_TO_CODE[matched[0]]

        if matched[1] == '同':
            i = int(self.board.last_move_txt[3]) - 1
            j = int(self.board.last_move_txt[4]) - 1
        else:
            i = TO_INT[matched[1][0]] - 1
            j = TO_INT[matched[1][1]] - 1

        piece = KANJI_TO_PIECE[matched[2]]

        teban = int(code == '+')
        direction = 2 * teban - 1

        pieces_index = self.board.get_piece_indexes(piece)[1 - teban]
        prev_pos_candidates = []

        # Search candidate pieces that can move to (i, j)
        for pi, pj in pieces_index:
            koma_act = PIECE_TO_ACT[piece]
            for act in koma_act:
                for move in act:

                    next_i = pi + move[0] * direction
                    next_j = pj + move[1] * direction

                    # If next_i or next_j is outside of the board
                    if (next_i < 0 or 9 <= next_i or
                            next_j < 0 or 9 <= next_j):
                        break

                    # If conflict with the piece of move_ki2
                    if i == next_i and j == next_j:
                        prev_pos_candidates.append([pi, pj])

                    # If conflict with other pieces
                    if self.board[next_i][next_j] != EMPTY_STR:
                        break

        # If promote
        if matched[3] and re.findall('(?<!不)成', matched[3]):
            piece = TURN_PIECE_REVERSED[piece]

        if matched[3] and '打' in matched[3]:
            prev_pos_candidates = [[-1, -1]]

        # Put piece from mochigoma
        if not prev_pos_candidates:
            prev_pos_candidates = [[-1, -1]]

        n_candidates = len(prev_pos_candidates)
        if n_candidates == 1:
            prev_pos = prev_pos_candidates[0]

        elif n_candidates >= 2:
            if matched[3]:
                if '右' in matched[3]:
                    if piece in OGOMA:
                        prev_pos_candidates = filter(
                            lambda x: (i - x[0]) * direction >= 0,
                            prev_pos_candidates)
                    else:
                        prev_pos_candidates = filter(
                            lambda x: (i - x[0]) * direction > 0,
                            prev_pos_candidates)
                elif '左' in matched[3]:
                    if piece in OGOMA:
                        prev_pos_candidates = filter(
                            lambda x: (i - x[0]) * direction <= 0,
                            prev_pos_candidates)
                    else:
                        prev_pos_candidates = filter(
                            lambda x: (i - x[0]) * direction < 0,
                            prev_pos_candidates)

                if '上' in matched[3] or '行' in matched[3]:
                    prev_pos_candidates = filter(
                        lambda x: (x[1] - j) * direction > 0,
                        prev_pos_candidates
                    )
                elif '直' in matched[3]:
                    prev_pos_candidates = filter(
                        lambda x: x[0] == i and (x[1] - j) * direction > 0,
                        prev_pos_candidates
                    )
                elif '寄' in matched[3]:
                    prev_pos_candidates = filter(
                        lambda x: x[1] == j,
                        prev_pos_candidates
                    )
                elif '引' in matched[3]:
                    prev_pos_candidates = filter(
                        lambda x: (x[1] - j) * direction < 0,
                        prev_pos_candidates
                    )

                # cf. http://www.shogi.or.jp/faq/kihuhyouki.html
                # "竜が2枚の場合はやはり動作を優先します。
                #  ただし、「直」は使わずに「左」「右」で記入します。"
                if (piece in OGOMA and
                    ('右' in matched[3] or '左' in matched[3])
                ):
                    if '右' in matched[3]:
                        lr = -1
                    elif '左' in matched[3]:
                        lr = 1

                    prev_pos_candidates = [
                        max(prev_pos_candidates,
                            key=lambda x: lr * direction * x[0])
                    ]

                prev_pos_candidates = list(prev_pos_candidates)
                if (not prev_pos_candidates or
                        len(prev_pos_candidates) >= 2):
                    raise RuntimeError('Parse Error')

                prev_pos = prev_pos_candidates[0]
        else:
            raise RuntimeError('Cannot find prev position.')

        if prev_pos and len(prev_pos) == 2:
            move_csa = ''.join([
                code,
                str(prev_pos[0] + 1),
                str(prev_pos[1] + 1),
                str(i + 1),
                str(j + 1),
                piece
            ])

            return move_csa

        else:
            return None
<|endoftext|>


<commit_before>#!/usr/bin/python
#
# Copyright (c) 2012 Mikkel Schubert <MSchubert@snm.ku.dk>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
import os
import sys
import signal
import types
import weakref
import subprocess
import collections

import pypeline.common.fileutils as fileutils

_PIPES = ('IN_STDIN', 'OUT_STDOUT', 'OUT_STDERR')
_PREFIXES = ('IN_', 'TEMP_IN_', 'OUT_', 'TEMP_OUT_', 'EXEC_', 'AUX_', 'CHECK_')

class CmdError(RuntimeError):
    def __init__(self, msg):
        RuntimeError.__init__(self, msg)

class AtomicCmd:
    """Executes a command, only moving resulting files to the destination
    directory if the command was succesful. This helps prevent the
    accidential use of partial files in downstream analysis, and eases
    restarting of a pipeline following errors (no cleanup).

    Inividual files/paths in the command are specified using keywords (see
    the documentation for the constructor), allowing the command to be
    transparently modified to execute in a temporary directory.

    When an AtomicCmd is run(), a signal handler is installed for SIGTERM,
    which ensures that any running processes are terminated. In the absence
    of this, AtomicCmds run in terminated subprocesses can result in still
    running children after the termination of the parents."""
    PIPE = subprocess.PIPE

    def __init__(self, command, set_cwd = False, **kwargs):
        """Takes a command and a set of files.

        The command is expected to be an iterable starting with the name of an
        executable, with each item representing one string on the command line.
        Thus, the command "find /etc -name 'profile*'" might be represented as
        the list ["find", "/etc", "-name", "profile*"].

        If 'set_cwd' is True, the current working directory is set to the
        temporary directory before the command is executed. Input paths are
        automatically turned into absolute paths in this case.

        Each keyword represents a type of file, as determined by the prefix:
           IN_    -- Path to input file transformed/analysed the executable.
           OUT_   -- Path to output file generated by the executable. During
                     execution of the AtomicCmd, this is modified to point to
                     the temporary directory.
           EXEC_  -- Name of / path to executable, or RequirementObj. The first
                     item in the command is always one of the executables, even
                     if not specified in this manner.
           AUX_   -- Auxillery files required by the executable(s), which are
                     themselves not executable. Examples include scripts,
                     config files, data-bases, and the like.
           CHECK_ -- A callable, which upon calling carries out version checking,
                     raising an exception in the case of requirements not being
                     met. This may be used to help ensure that prerequisites are
                     met before running the command.

        Note that files that are not directly invoked may be included above,
        in order to allow the specification of requirements. This could include
        required data files, or executables indirectly executed by a script.

        If the above is prefixed with "TEMP_", the files are read from / written
        to the temporary folder in which the command is executed. Note that all
        TEMP_OUT_ files are deleted when commit is called (if they exist).

        In addition, the follow special names may be used with the above:
           STDIN_  -- Takes a filename, or an AtomicCmd, in which case the stdout
                      of that command is piped to the stdin of this instance.
           STDOUT_ -- Takes a filename, or the special value PIPE to allow
                      another AtomicCmd instance to use the output directly.
           STDERR_ -- Takes a filename.

        Each pipe can only be used once (either OUT_ or TEMP_OUT_)."""

        self._proc    = None
        self._command = [str(field) for field in command]
        self._handles = []
        self._set_cwd = set_cwd

        self._files     = self._process_arguments(id(self), command, kwargs)
        self._file_sets = self._build_files_map(command, kwargs)

    def run(self, temp):
        """Runs the given command, saving files in the specified temp folder. To
        move files to their final destination, call commit(). Note that in contexts
        where the *Cmds classes are used, this function may block."""

        # kwords for pipes are always built relative to the current directory,
        # since these are opened before (possibly) CD'ing to the temp directory.
        kwords = self._generate_filenames(self._files, root = temp)
        stdin  = self._open_pipe(kwords, "IN_STDIN" , "rb")
        stdout = self._open_pipe(kwords, "OUT_STDOUT", "wb")
        stderr = self._open_pipe(kwords, "OUT_STDERR", "wb")

        cwd = None
        if self._set_cwd:
            cwd, kwords = temp, self._generate_filenames(self._files, root = None)

        command = [(field % kwords) for field in self._command]
        self._proc = subprocess.Popen(command,
                                      stdin  = stdin,
                                      stdout = stdout,
                                      stderr = stderr,
                                      cwd    = cwd,
                                      preexec_fn = os.setsid,
                                      close_fds = True)

        # Allow subprocesses to be killed in case of a SIGTERM
        _add_to_killlist(self._proc)

    def ready(self):
        """Returns true if the command has been run to completion,
        regardless of wether or not an error occured."""
        return (self._proc and self._proc.poll() is not None)

    def join(self):
        """Similar to Popen.wait(), but returns the value wrapped in a list,
        and ensures that any opened handles are closed. Must be called before
        calling commit."""
        try:
            return_codes = [self._proc.wait()] if self._proc else [None]
        finally:
            # Close any implictly opened pipes
            for (mode, handle) in self._handles:
                if "w" in mode:
                    handle.flush()
                handle.close()
            self._handles = []

        return return_codes

    def wait(self):
        """Equivalent to Subproces.wait. This function should only
        be used in contexts where a AtomicCmd needs to be combined
        with Subprocesses, as it does not exist for AtomicSets."""
        return self.join()[0]

    def terminate(self):
        """Sends SIGTERM to process."""
        self._proc.terminate()
        self._proc = None

    @property
    def executables(self):
        """Returns a list of executables required for the AtomicCmd."""
        return self._file_sets["executable"]

    @property
    def requirements(self):
        """Returns a list of callable requirements that must be met
        before the AtomicCmd can be expected to be runable."""
        return self._file_sets["requirements"]

    @property
    def input_files(self):
        """Returns a list of input files that are required by the AtomicCmd."""
        return self._file_sets["input"]

    @property
    def output_files(self):
        """Checks that the expected output files have been generated."""
        return self._file_sets["output"]

    @property
    def auxiliary_files(self):
        """Checks that the expected output files have been generated."""
        return self._file_sets["auxiliary"]

    def commit(self, temp):
        if not self.ready():
            raise CmdError("Attempting to commit command before it has completed")
        elif self._handles:
            raise CmdError("Called 'commit' before calling 'join'")

        self._proc = None

        temp_files = self._generate_filenames(self._files, temp)
        for (key, filename) in temp_files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("OUT_") and not os.path.exists(filename):
                    raise CmdError("Command did not create expected output file: " + filename)

        for (key, filename) in temp_files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("OUT_"):
                    fileutils.move_file(filename, self._files[key])
                elif key.startswith("TEMP_OUT_") and os.path.exists(filename):
                    os.remove(filename)

    @property
    def stdout(self):
        """Returns the 'stdout' value for the current process. If no
        such process has been started, the return value is None."""
        if self._proc:
            return self._proc.stdout

    def __str__(self):
        def describe_pipe(template, pipe):
            if isinstance(pipe, types.StringTypes):
                return template % pipe
            elif isinstance(pipe, AtomicCmd):
                return template % "[AtomicCmd]"
            elif (pipe == AtomicCmd.PIPE):
                return template % "[PIPE]"
            else:
                return ""

        kwords = self._generate_filenames(self._files, "${TEMP}")
        command = " ".join([(field % kwords) for field in self._command])

        stdin = self._files.get("IN_STDIN")
        if stdin:
            command += describe_pipe(" < %s", stdin)

        stdout = self._files.get("OUT_STDOUT")
        stderr = self._files.get("OUT_STDERR")
        if (stdout != stderr):
            if stdout:
                command += describe_pipe(" > %s", stdout)
            if stderr:
                command += describe_pipe(" 2> %s", stderr)
        elif stdout:
            command += describe_pipe(" &> %s", stdout)

        return "<%s>" % command

    @classmethod
    def _process_arguments(cls, proc_id, command, kwargs):
        cls._validate_pipes(kwargs)

        files = {}
        for (key, value) in kwargs.iteritems():
            if cls._validate_argument(key, value):
                files[key] = value

        executable = os.path.basename(command[0])
        for pipe in ("STDOUT", "STDERR"):
            if not (kwargs.get("OUT_" + pipe) or kwargs.get("TEMP_OUT_" + pipe)):
                filename = "pipe_%s_%i.%s" % (executable, proc_id, pipe.lower())
                files["TEMP_OUT_" + pipe] = filename

        output_files = collections.defaultdict(list)
        for (key, filename) in kwargs.iteritems():
            if key.startswith("TEMP_OUT_") or key.startswith("OUT_"):
                if isinstance(filename, types.StringTypes):
                    output_files[os.path.basename(filename)].append(key)

        for (filename, keys) in output_files.iteritems():
            if len(keys) > 1:
                raise CmdError("Same output filename (%s) is specified for multiple keys: %s" \
                                   % (filename, ", ".join(keys)))

        return files

    @classmethod
    def _validate_pipes(cls, kwargs):
        """Checks that no single pipe is specified multiple times, e.i. being specified
        both for a temporary and a final (outside the temp dir) file. For example,
        either IN_STDIN or TEMP_IN_STDIN must be specified, but not both."""
        if any((kwargs.get(pipe) and kwargs.get("TEMP_" + pipe)) for pipe in _PIPES):
            raise CmdError, "Pipes must be specified at most once (w/wo TEMP_)."

    @classmethod
    def _validate_argument(cls, key, value):
        if not any(key.startswith(prefix) for prefix in _PREFIXES):
            raise CmdError("Command contains invalid argument (wrong prefix): '%s' -> '%s'" \
                           % (cls.__name__, key))
        elif (value == cls.PIPE) and key not in ("OUT_STDOUT", "TEMP_OUT_STDOUT"):
            raise ValueError, "PIPE is only allow for *_STDOUT, not " + key
        elif key.startswith("CHECK_") and not isinstance(value, collections.Callable):
            raise ValueError("Values CHECKS_ must be callables, not %s" % repr(value))
        elif (key == "IN_STDIN") and not isinstance(value, types.StringTypes + (AtomicCmd,)):
            raise ValueError("STDIN for '%s' is not a string or AtomicCmd: %s" \
                             % (cls.__name__, value))

        # Values are allowed to be none, but such entries are skipped
        # This simplifies the use of keyword parameters with no default values.
        return bool(value)

    def _open_pipe(self, kwords, pipe, mode):
        filename = kwords.get(pipe, kwords.get("TEMP_" + pipe))
        if filename in (None, self.PIPE):
            return filename
        elif isinstance(filename, AtomicCmd):
            return filename.stdout

        pipe = open(filename, mode)
        self._handles.append((mode, pipe))

        return pipe

    @classmethod
    def _generate_filenames(cls, files, root):
        root = os.path.abspath(root) if root else ""

        filenames = {"TEMP_DIR" : root}
        for (key, filename) in files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("TEMP_") or key.startswith("OUT_"):
                    filename = os.path.join(root, os.path.basename(filename))
                elif not root and (key.startswith("IN_") or key.startswith("AUX_")):
                    filename = os.path.abspath(filename)
            filenames[key] = filename

        return filenames

    @classmethod
    def _build_files_map(cls, command, files):
        key_map   = {"IN" : "input",  "OUT": "output",  "EXEC" : "executable",
                     "AUX" : "auxiliary", "CHECK" : "requirements"}
        file_sets = dict((key, set()) for key in key_map.itervalues())

        file_sets["executable"].add(command[0])
        for (key, filename) in files.iteritems():
            if not isinstance(filename, types.StringTypes) and not key.startswith("CHECK_"):
                continue
            elif not key.startswith("TEMP_"):
                key = key_map[key.split("_", 1)[0]]
                file_sets[key].add(filename)

        return dict(zip(file_sets.keys(), map(frozenset, file_sets.values())))

## The following ensures proper cleanup of child processes, for example in the case
## where multiprocessing.Pool.terminate() is called. The current implementation will
## leak weakref objects, but given the low number of commands called over the lifetime
## of a pipeline, this is considered acceptable for now. FIXME
_PROCS = []

def _cleanup_children(signum, _frame):
    for proc_ref in _PROCS:
        proc = proc_ref()
        if proc:
            os.killpg(proc.pid, signal.SIGTERM)
    sys.exit(-signum)

def _add_to_killlist(proc):
    if not _PROCS:
        signal.signal(signal.SIGTERM, _cleanup_children)

    _PROCS.append(weakref.ref(proc))
<commit_msg>Ignore termination of non-running/already terminated cmds.<commit_after>#!/usr/bin/python
#
# Copyright (c) 2012 Mikkel Schubert <MSchubert@snm.ku.dk>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
import os
import sys
import signal
import types
import weakref
import subprocess
import collections

import pypeline.common.fileutils as fileutils

_PIPES = ('IN_STDIN', 'OUT_STDOUT', 'OUT_STDERR')
_PREFIXES = ('IN_', 'TEMP_IN_', 'OUT_', 'TEMP_OUT_', 'EXEC_', 'AUX_', 'CHECK_')

class CmdError(RuntimeError):
    def __init__(self, msg):
        RuntimeError.__init__(self, msg)

class AtomicCmd:
    """Executes a command, only moving resulting files to the destination
    directory if the command was succesful. This helps prevent the
    accidential use of partial files in downstream analysis, and eases
    restarting of a pipeline following errors (no cleanup).

    Inividual files/paths in the command are specified using keywords (see
    the documentation for the constructor), allowing the command to be
    transparently modified to execute in a temporary directory.

    When an AtomicCmd is run(), a signal handler is installed for SIGTERM,
    which ensures that any running processes are terminated. In the absence
    of this, AtomicCmds run in terminated subprocesses can result in still
    running children after the termination of the parents."""
    PIPE = subprocess.PIPE

    def __init__(self, command, set_cwd = False, **kwargs):
        """Takes a command and a set of files.

        The command is expected to be an iterable starting with the name of an
        executable, with each item representing one string on the command line.
        Thus, the command "find /etc -name 'profile*'" might be represented as
        the list ["find", "/etc", "-name", "profile*"].

        If 'set_cwd' is True, the current working directory is set to the
        temporary directory before the command is executed. Input paths are
        automatically turned into absolute paths in this case.

        Each keyword represents a type of file, as determined by the prefix:
           IN_    -- Path to input file transformed/analysed the executable.
           OUT_   -- Path to output file generated by the executable. During
                     execution of the AtomicCmd, this is modified to point to
                     the temporary directory.
           EXEC_  -- Name of / path to executable, or RequirementObj. The first
                     item in the command is always one of the executables, even
                     if not specified in this manner.
           AUX_   -- Auxillery files required by the executable(s), which are
                     themselves not executable. Examples include scripts,
                     config files, data-bases, and the like.
           CHECK_ -- A callable, which upon calling carries out version checking,
                     raising an exception in the case of requirements not being
                     met. This may be used to help ensure that prerequisites are
                     met before running the command.

        Note that files that are not directly invoked may be included above,
        in order to allow the specification of requirements. This could include
        required data files, or executables indirectly executed by a script.

        If the above is prefixed with "TEMP_", the files are read from / written
        to the temporary folder in which the command is executed. Note that all
        TEMP_OUT_ files are deleted when commit is called (if they exist).

        In addition, the follow special names may be used with the above:
           STDIN_  -- Takes a filename, or an AtomicCmd, in which case the stdout
                      of that command is piped to the stdin of this instance.
           STDOUT_ -- Takes a filename, or the special value PIPE to allow
                      another AtomicCmd instance to use the output directly.
           STDERR_ -- Takes a filename.

        Each pipe can only be used once (either OUT_ or TEMP_OUT_)."""

        self._proc    = None
        self._command = [str(field) for field in command]
        self._handles = []
        self._set_cwd = set_cwd

        self._files     = self._process_arguments(id(self), command, kwargs)
        self._file_sets = self._build_files_map(command, kwargs)

    def run(self, temp):
        """Runs the given command, saving files in the specified temp folder. To
        move files to their final destination, call commit(). Note that in contexts
        where the *Cmds classes are used, this function may block."""

        # kwords for pipes are always built relative to the current directory,
        # since these are opened before (possibly) CD'ing to the temp directory.
        kwords = self._generate_filenames(self._files, root = temp)
        stdin  = self._open_pipe(kwords, "IN_STDIN" , "rb")
        stdout = self._open_pipe(kwords, "OUT_STDOUT", "wb")
        stderr = self._open_pipe(kwords, "OUT_STDERR", "wb")

        cwd = None
        if self._set_cwd:
            cwd, kwords = temp, self._generate_filenames(self._files, root = None)

        command = [(field % kwords) for field in self._command]
        self._proc = subprocess.Popen(command,
                                      stdin  = stdin,
                                      stdout = stdout,
                                      stderr = stderr,
                                      cwd    = cwd,
                                      preexec_fn = os.setsid,
                                      close_fds = True)

        # Allow subprocesses to be killed in case of a SIGTERM
        _add_to_killlist(self._proc)

    def ready(self):
        """Returns true if the command has been run to completion,
        regardless of wether or not an error occured."""
        return (self._proc and self._proc.poll() is not None)

    def join(self):
        """Similar to Popen.wait(), but returns the value wrapped in a list,
        and ensures that any opened handles are closed. Must be called before
        calling commit."""
        try:
            return_codes = [self._proc.wait()] if self._proc else [None]
        finally:
            # Close any implictly opened pipes
            for (mode, handle) in self._handles:
                if "w" in mode:
                    handle.flush()
                handle.close()
            self._handles = []

        return return_codes

    def wait(self):
        """Equivalent to Subproces.wait. This function should only
        be used in contexts where a AtomicCmd needs to be combined
        with Subprocesses, as it does not exist for AtomicSets."""
        return self.join()[0]

    def terminate(self):
        """Sends SIGTERM to process."""
        if self._proc:
            self._proc.terminate()
            self._proc = None

    @property
    def executables(self):
        """Returns a list of executables required for the AtomicCmd."""
        return self._file_sets["executable"]

    @property
    def requirements(self):
        """Returns a list of callable requirements that must be met
        before the AtomicCmd can be expected to be runable."""
        return self._file_sets["requirements"]

    @property
    def input_files(self):
        """Returns a list of input files that are required by the AtomicCmd."""
        return self._file_sets["input"]

    @property
    def output_files(self):
        """Checks that the expected output files have been generated."""
        return self._file_sets["output"]

    @property
    def auxiliary_files(self):
        """Checks that the expected output files have been generated."""
        return self._file_sets["auxiliary"]

    def commit(self, temp):
        if not self.ready():
            raise CmdError("Attempting to commit command before it has completed")
        elif self._handles:
            raise CmdError("Called 'commit' before calling 'join'")

        self._proc = None

        temp_files = self._generate_filenames(self._files, temp)
        for (key, filename) in temp_files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("OUT_") and not os.path.exists(filename):
                    raise CmdError("Command did not create expected output file: " + filename)

        for (key, filename) in temp_files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("OUT_"):
                    fileutils.move_file(filename, self._files[key])
                elif key.startswith("TEMP_OUT_") and os.path.exists(filename):
                    os.remove(filename)

    @property
    def stdout(self):
        """Returns the 'stdout' value for the current process. If no
        such process has been started, the return value is None."""
        if self._proc:
            return self._proc.stdout

    def __str__(self):
        def describe_pipe(template, pipe):
            if isinstance(pipe, types.StringTypes):
                return template % pipe
            elif isinstance(pipe, AtomicCmd):
                return template % "[AtomicCmd]"
            elif (pipe == AtomicCmd.PIPE):
                return template % "[PIPE]"
            else:
                return ""

        kwords = self._generate_filenames(self._files, "${TEMP}")
        command = " ".join([(field % kwords) for field in self._command])

        stdin = self._files.get("IN_STDIN")
        if stdin:
            command += describe_pipe(" < %s", stdin)

        stdout = self._files.get("OUT_STDOUT")
        stderr = self._files.get("OUT_STDERR")
        if (stdout != stderr):
            if stdout:
                command += describe_pipe(" > %s", stdout)
            if stderr:
                command += describe_pipe(" 2> %s", stderr)
        elif stdout:
            command += describe_pipe(" &> %s", stdout)

        return "<%s>" % command

    @classmethod
    def _process_arguments(cls, proc_id, command, kwargs):
        cls._validate_pipes(kwargs)

        files = {}
        for (key, value) in kwargs.iteritems():
            if cls._validate_argument(key, value):
                files[key] = value

        executable = os.path.basename(command[0])
        for pipe in ("STDOUT", "STDERR"):
            if not (kwargs.get("OUT_" + pipe) or kwargs.get("TEMP_OUT_" + pipe)):
                filename = "pipe_%s_%i.%s" % (executable, proc_id, pipe.lower())
                files["TEMP_OUT_" + pipe] = filename

        output_files = collections.defaultdict(list)
        for (key, filename) in kwargs.iteritems():
            if key.startswith("TEMP_OUT_") or key.startswith("OUT_"):
                if isinstance(filename, types.StringTypes):
                    output_files[os.path.basename(filename)].append(key)

        for (filename, keys) in output_files.iteritems():
            if len(keys) > 1:
                raise CmdError("Same output filename (%s) is specified for multiple keys: %s" \
                                   % (filename, ", ".join(keys)))

        return files

    @classmethod
    def _validate_pipes(cls, kwargs):
        """Checks that no single pipe is specified multiple times, e.i. being specified
        both for a temporary and a final (outside the temp dir) file. For example,
        either IN_STDIN or TEMP_IN_STDIN must be specified, but not both."""
        if any((kwargs.get(pipe) and kwargs.get("TEMP_" + pipe)) for pipe in _PIPES):
            raise CmdError, "Pipes must be specified at most once (w/wo TEMP_)."

    @classmethod
    def _validate_argument(cls, key, value):
        if not any(key.startswith(prefix) for prefix in _PREFIXES):
            raise CmdError("Command contains invalid argument (wrong prefix): '%s' -> '%s'" \
                           % (cls.__name__, key))
        elif (value == cls.PIPE) and key not in ("OUT_STDOUT", "TEMP_OUT_STDOUT"):
            raise ValueError, "PIPE is only allow for *_STDOUT, not " + key
        elif key.startswith("CHECK_") and not isinstance(value, collections.Callable):
            raise ValueError("Values CHECKS_ must be callables, not %s" % repr(value))
        elif (key == "IN_STDIN") and not isinstance(value, types.StringTypes + (AtomicCmd,)):
            raise ValueError("STDIN for '%s' is not a string or AtomicCmd: %s" \
                             % (cls.__name__, value))

        # Values are allowed to be none, but such entries are skipped
        # This simplifies the use of keyword parameters with no default values.
        return bool(value)

    def _open_pipe(self, kwords, pipe, mode):
        filename = kwords.get(pipe, kwords.get("TEMP_" + pipe))
        if filename in (None, self.PIPE):
            return filename
        elif isinstance(filename, AtomicCmd):
            return filename.stdout

        pipe = open(filename, mode)
        self._handles.append((mode, pipe))

        return pipe

    @classmethod
    def _generate_filenames(cls, files, root):
        root = os.path.abspath(root) if root else ""

        filenames = {"TEMP_DIR" : root}
        for (key, filename) in files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("TEMP_") or key.startswith("OUT_"):
                    filename = os.path.join(root, os.path.basename(filename))
                elif not root and (key.startswith("IN_") or key.startswith("AUX_")):
                    filename = os.path.abspath(filename)
            filenames[key] = filename

        return filenames

    @classmethod
    def _build_files_map(cls, command, files):
        key_map   = {"IN" : "input",  "OUT": "output",  "EXEC" : "executable",
                     "AUX" : "auxiliary", "CHECK" : "requirements"}
        file_sets = dict((key, set()) for key in key_map.itervalues())

        file_sets["executable"].add(command[0])
        for (key, filename) in files.iteritems():
            if not isinstance(filename, types.StringTypes) and not key.startswith("CHECK_"):
                continue
            elif not key.startswith("TEMP_"):
                key = key_map[key.split("_", 1)[0]]
                file_sets[key].add(filename)

        return dict(zip(file_sets.keys(), map(frozenset, file_sets.values())))

## The following ensures proper cleanup of child processes, for example in the case
## where multiprocessing.Pool.terminate() is called. The current implementation will
## leak weakref objects, but given the low number of commands called over the lifetime
## of a pipeline, this is considered acceptable for now. FIXME
_PROCS = []

def _cleanup_children(signum, _frame):
    for proc_ref in _PROCS:
        proc = proc_ref()
        if proc:
            os.killpg(proc.pid, signal.SIGTERM)
    sys.exit(-signum)

def _add_to_killlist(proc):
    if not _PROCS:
        signal.signal(signal.SIGTERM, _cleanup_children)

    _PROCS.append(weakref.ref(proc))
<|endoftext|>


<commit_before>#!/usr/bin/python
#
# Copyright (c) 2012 Mikkel Schubert <MSchubert@snm.ku.dk>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
import os
import sys
import signal
import types
import weakref
import subprocess
import collections

import pypeline.common.fileutils as fileutils

_PIPES = ('IN_STDIN', 'OUT_STDOUT', 'OUT_STDERR')
_PREFIXES = ('IN_', 'TEMP_IN_', 'OUT_', 'TEMP_OUT_', 'EXEC_', 'AUX_', 'CHECK_')

class CmdError(RuntimeError):
    def __init__(self, msg):
        RuntimeError.__init__(self, msg)

class AtomicCmd:
    """Executes a command, only moving resulting files to the destination
    directory if the command was succesful. This helps prevent the
    accidential use of partial files in downstream analysis, and eases
    restarting of a pipeline following errors (no cleanup).

    Inividual files/paths in the command are specified using keywords (see
    the documentation for the constructor), allowing the command to be
    transparently modified to execute in a temporary directory.

    When an AtomicCmd is run(), a signal handler is installed for SIGTERM,
    which ensures that any running processes are terminated. In the absence
    of this, AtomicCmds run in terminated subprocesses can result in still
    running children after the termination of the parents."""
    PIPE = subprocess.PIPE

    def __init__(self, command, set_cwd = False, **kwargs):
        """Takes a command and a set of files.

        The command is expected to be an iterable starting with the name of an
        executable, with each item representing one string on the command line.
        Thus, the command "find /etc -name 'profile*'" might be represented as
        the list ["find", "/etc", "-name", "profile*"].

        If 'set_cwd' is True, the current working directory is set to the
        temporary directory before the command is executed. Input paths are
        automatically turned into absolute paths in this case.

        Each keyword represents a type of file, as determined by the prefix:
           IN_    -- Path to input file transformed/analysed the executable.
           OUT_   -- Path to output file generated by the executable. During
                     execution of the AtomicCmd, this is modified to point to
                     the temporary directory.
           EXEC_  -- Name of / path to executable, or RequirementObj. The first
                     item in the command is always one of the executables, even
                     if not specified in this manner.
           AUX_   -- Auxillery files required by the executable(s), which are
                     themselves not executable. Examples include scripts,
                     config files, data-bases, and the like.
           CHECK_ -- A callable, which upon calling carries out version checking,
                     raising an exception in the case of requirements not being
                     met. This may be used to help ensure that prerequisites are
                     met before running the command.

        Note that files that are not directly invoked may be included above,
        in order to allow the specification of requirements. This could include
        required data files, or executables indirectly executed by a script.

        If the above is prefixed with "TEMP_", the files are read from / written
        to the temporary folder in which the command is executed. Note that all
        TEMP_OUT_ files are deleted when commit is called (if they exist).

        In addition, the follow special names may be used with the above:
           STDIN_  -- Takes a filename, or an AtomicCmd, in which case the stdout
                      of that command is piped to the stdin of this instance.
           STDOUT_ -- Takes a filename, or the special value PIPE to allow
                      another AtomicCmd instance to use the output directly.
           STDERR_ -- Takes a filename.

        Each pipe can only be used once (either OUT_ or TEMP_OUT_)."""

        self._proc    = None
        self._command = [str(field) for field in command]
        self._handles = []
        self._set_cwd = set_cwd

        self._files     = self._process_arguments(id(self), command, kwargs)
        self._file_sets = self._build_files_map(command, kwargs)

    def run(self, temp):
        """Runs the given command, saving files in the specified temp folder. To
        move files to their final destination, call commit(). Note that in contexts
        where the *Cmds classes are used, this function may block."""

        # kwords for pipes are always built relative to the current directory,
        # since these are opened before (possibly) CD'ing to the temp directory.
        kwords = self._generate_filenames(self._files, root = temp)
        stdin  = self._open_pipe(kwords, "IN_STDIN" , "rb")
        stdout = self._open_pipe(kwords, "OUT_STDOUT", "wb")
        stderr = self._open_pipe(kwords, "OUT_STDERR", "wb")

        cwd = None
        if self._set_cwd:
            cwd, kwords = temp, self._generate_filenames(self._files, root = None)

        command = [(field % kwords) for field in self._command]
        self._proc = subprocess.Popen(command,
                                      stdin  = stdin,
                                      stdout = stdout,
                                      stderr = stderr,
                                      cwd    = cwd,
                                      close_fds = True)

        # Allow subprocesses to be killed in case of a SIGTERM
        _add_to_killlist(self._proc)

    def ready(self):
        """Returns true if the command has been run to completion,
        regardless of wether or not an error occured."""
        return (self._proc and self._proc.poll() is not None)

    def join(self):
        """Similar to Popen.wait(), but returns the value wrapped in a list,
        and ensures that any opened handles are closed. Must be called before
        calling commit."""
        try:
            return_codes = [self._proc.wait()] if self._proc else [None]
        finally:
            # Close any implictly opened pipes
            for (mode, handle) in self._handles:
                if "w" in mode:
                    handle.flush()
                handle.close()
            self._handles = []

        return return_codes

    def wait(self):
        """Equivalent to Subproces.wait. This function should only
        be used in contexts where a AtomicCmd needs to be combined
        with Subprocesses, as it does not exist for AtomicSets."""
        return self.join()[0]

    def terminate(self):
        """Sends SIGTERM to process."""
        self._proc.terminate()
        self._proc = None

    @property
    def executables(self):
        """Returns a list of executables required for the AtomicCmd."""
        return self._file_sets["executable"]

    @property
    def requirements(self):
        """Returns a list of callable requirements that must be met
        before the AtomicCmd can be expected to be runable."""
        return self._file_sets["requirements"]

    @property
    def input_files(self):
        """Returns a list of input files that are required by the AtomicCmd."""
        return self._file_sets["input"]

    @property
    def output_files(self):
        """Checks that the expected output files have been generated."""
        return self._file_sets["output"]

    @property
    def auxiliary_files(self):
        """Checks that the expected output files have been generated."""
        return self._file_sets["auxiliary"]

    def commit(self, temp):
        if not self.ready():
            raise CmdError("Attempting to commit command before it has completed")
        elif self._handles:
            raise CmdError("Called 'commit' before calling 'join'")

        self._proc = None

        temp_files = self._generate_filenames(self._files, temp)
        for (key, filename) in temp_files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("OUT_") and not os.path.exists(filename):
                    raise CmdError("Command did not create expected output file: " + filename)

        for (key, filename) in temp_files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("OUT_"):
                    fileutils.move_file(filename, self._files[key])
                elif key.startswith("TEMP_OUT_") and os.path.exists(filename):
                    os.remove(filename)

    @property
    def stdout(self):
        """Returns the 'stdout' value for the current process. If no
        such process has been started, the return value is None."""
        if self._proc:
            return self._proc.stdout

    def __str__(self):
        def describe_pipe(template, pipe):
            if isinstance(pipe, types.StringTypes):
                return template % pipe
            elif isinstance(pipe, AtomicCmd):
                return template % "[AtomicCmd]"
            elif (pipe == AtomicCmd.PIPE):
                return template % "[PIPE]"
            else:
                return ""

        kwords = self._generate_filenames(self._files, "${TEMP}")
        command = " ".join([(field % kwords) for field in self._command])

        stdin = self._files.get("IN_STDIN")
        if stdin:
            command += describe_pipe(" < %s", stdin)

        stdout = self._files.get("OUT_STDOUT")
        stderr = self._files.get("OUT_STDERR")
        if (stdout != stderr):
            if stdout:
                command += describe_pipe(" > %s", stdout)
            if stderr:
                command += describe_pipe(" 2> %s", stderr)
        elif stdout:
            command += describe_pipe(" &> %s", stdout)

        return "<%s>" % command

    @classmethod
    def _process_arguments(cls, proc_id, command, kwargs):
        cls._validate_pipes(kwargs)

        files = {}
        for (key, value) in kwargs.iteritems():
            if cls._validate_argument(key, value):
                files[key] = value

        executable = os.path.basename(command[0])
        for pipe in ("STDOUT", "STDERR"):
            if not (kwargs.get("OUT_" + pipe) or kwargs.get("TEMP_OUT_" + pipe)):
                filename = "pipe_%s_%i.%s" % (executable, proc_id, pipe.lower())
                files["TEMP_OUT_" + pipe] = filename

        output_files = collections.defaultdict(list)
        for (key, filename) in kwargs.iteritems():
            if key.startswith("TEMP_OUT_") or key.startswith("OUT_"):
                if isinstance(filename, types.StringTypes):
                    output_files[os.path.basename(filename)].append(key)

        for (filename, keys) in output_files.iteritems():
            if len(keys) > 1:
                raise CmdError("Same output filename (%s) is specified for multiple keys: %s" \
                                   % (filename, ", ".join(keys)))

        return files

    @classmethod
    def _validate_pipes(cls, kwargs):
        """Checks that no single pipe is specified multiple times, e.i. being specified
        both for a temporary and a final (outside the temp dir) file. For example,
        either IN_STDIN or TEMP_IN_STDIN must be specified, but not both."""
        if any((kwargs.get(pipe) and kwargs.get("TEMP_" + pipe)) for pipe in _PIPES):
            raise CmdError, "Pipes must be specified at most once (w/wo TEMP_)."

    @classmethod
    def _validate_argument(cls, key, value):
        if not any(key.startswith(prefix) for prefix in _PREFIXES):
            raise CmdError("Command contains invalid argument (wrong prefix): '%s' -> '%s'" \
                           % (cls.__name__, key))
        elif (value == cls.PIPE) and key not in ("OUT_STDOUT", "TEMP_OUT_STDOUT"):
            raise ValueError, "PIPE is only allow for *_STDOUT, not " + key
        elif key.startswith("CHECK_") and not isinstance(value, collections.Callable):
            raise ValueError("Values CHECKS_ must be callables, not %s" % repr(value))
        elif (key == "IN_STDIN") and not isinstance(value, types.StringTypes + (AtomicCmd,)):
            raise ValueError("STDIN for '%s' is not a string or AtomicCmd: %s" \
                             % (cls.__name__, value))

        # Values are allowed to be none, but such entries are skipped
        # This simplifies the use of keyword parameters with no default values.
        return bool(value)

    def _open_pipe(self, kwords, pipe, mode):
        filename = kwords.get(pipe, kwords.get("TEMP_" + pipe))
        if filename in (None, self.PIPE):
            return filename
        elif isinstance(filename, AtomicCmd):
            return filename.stdout

        pipe = open(filename, mode)
        self._handles.append((mode, pipe))

        return pipe

    @classmethod
    def _generate_filenames(cls, files, root):
        root = os.path.abspath(root) if root else ""

        filenames = {"TEMP_DIR" : root}
        for (key, filename) in files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("TEMP_") or key.startswith("OUT_"):
                    filename = os.path.join(root, os.path.basename(filename))
                elif not root and (key.startswith("IN_") or key.startswith("AUX_")):
                    filename = os.path.abspath(filename)
            filenames[key] = filename

        return filenames

    @classmethod
    def _build_files_map(cls, command, files):
        key_map   = {"IN" : "input",  "OUT": "output",  "EXEC" : "executable",
                     "AUX" : "auxiliary", "CHECK" : "requirements"}
        file_sets = dict((key, set()) for key in key_map.itervalues())

        file_sets["executable"].add(command[0])
        for (key, filename) in files.iteritems():
            if not isinstance(filename, types.StringTypes) and not key.startswith("CHECK_"):
                continue
            elif not key.startswith("TEMP_"):
                key = key_map[key.split("_", 1)[0]]
                file_sets[key].add(filename)

        return dict(zip(file_sets.keys(), map(frozenset, file_sets.values())))

## The following ensures proper cleanup of child processes, for example in the case
## where multiprocessing.Pool.terminate() is called. The current implementation will
## leak weakref objects, but given the low number of commands called over the lifetime
## of a pipeline, this is considered acceptable for now. FIXME
_PROCS = []

def _cleanup_children(signum, _frame):
    for proc_ref in _PROCS:
        proc = proc_ref()
        if proc:
            proc.terminate()
    sys.exit(-signum)

def _add_to_killlist(proc):
    if not _PROCS:
        signal.signal(signal.SIGTERM, _cleanup_children)

    _PROCS.append(weakref.ref(proc))
<commit_msg>Ensure that processes spawned by subprocesses are killed.<commit_after>#!/usr/bin/python
#
# Copyright (c) 2012 Mikkel Schubert <MSchubert@snm.ku.dk>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
import os
import sys
import signal
import types
import weakref
import subprocess
import collections

import pypeline.common.fileutils as fileutils

_PIPES = ('IN_STDIN', 'OUT_STDOUT', 'OUT_STDERR')
_PREFIXES = ('IN_', 'TEMP_IN_', 'OUT_', 'TEMP_OUT_', 'EXEC_', 'AUX_', 'CHECK_')

class CmdError(RuntimeError):
    def __init__(self, msg):
        RuntimeError.__init__(self, msg)

class AtomicCmd:
    """Executes a command, only moving resulting files to the destination
    directory if the command was succesful. This helps prevent the
    accidential use of partial files in downstream analysis, and eases
    restarting of a pipeline following errors (no cleanup).

    Inividual files/paths in the command are specified using keywords (see
    the documentation for the constructor), allowing the command to be
    transparently modified to execute in a temporary directory.

    When an AtomicCmd is run(), a signal handler is installed for SIGTERM,
    which ensures that any running processes are terminated. In the absence
    of this, AtomicCmds run in terminated subprocesses can result in still
    running children after the termination of the parents."""
    PIPE = subprocess.PIPE

    def __init__(self, command, set_cwd = False, **kwargs):
        """Takes a command and a set of files.

        The command is expected to be an iterable starting with the name of an
        executable, with each item representing one string on the command line.
        Thus, the command "find /etc -name 'profile*'" might be represented as
        the list ["find", "/etc", "-name", "profile*"].

        If 'set_cwd' is True, the current working directory is set to the
        temporary directory before the command is executed. Input paths are
        automatically turned into absolute paths in this case.

        Each keyword represents a type of file, as determined by the prefix:
           IN_    -- Path to input file transformed/analysed the executable.
           OUT_   -- Path to output file generated by the executable. During
                     execution of the AtomicCmd, this is modified to point to
                     the temporary directory.
           EXEC_  -- Name of / path to executable, or RequirementObj. The first
                     item in the command is always one of the executables, even
                     if not specified in this manner.
           AUX_   -- Auxillery files required by the executable(s), which are
                     themselves not executable. Examples include scripts,
                     config files, data-bases, and the like.
           CHECK_ -- A callable, which upon calling carries out version checking,
                     raising an exception in the case of requirements not being
                     met. This may be used to help ensure that prerequisites are
                     met before running the command.

        Note that files that are not directly invoked may be included above,
        in order to allow the specification of requirements. This could include
        required data files, or executables indirectly executed by a script.

        If the above is prefixed with "TEMP_", the files are read from / written
        to the temporary folder in which the command is executed. Note that all
        TEMP_OUT_ files are deleted when commit is called (if they exist).

        In addition, the follow special names may be used with the above:
           STDIN_  -- Takes a filename, or an AtomicCmd, in which case the stdout
                      of that command is piped to the stdin of this instance.
           STDOUT_ -- Takes a filename, or the special value PIPE to allow
                      another AtomicCmd instance to use the output directly.
           STDERR_ -- Takes a filename.

        Each pipe can only be used once (either OUT_ or TEMP_OUT_)."""

        self._proc    = None
        self._command = [str(field) for field in command]
        self._handles = []
        self._set_cwd = set_cwd

        self._files     = self._process_arguments(id(self), command, kwargs)
        self._file_sets = self._build_files_map(command, kwargs)

    def run(self, temp):
        """Runs the given command, saving files in the specified temp folder. To
        move files to their final destination, call commit(). Note that in contexts
        where the *Cmds classes are used, this function may block."""

        # kwords for pipes are always built relative to the current directory,
        # since these are opened before (possibly) CD'ing to the temp directory.
        kwords = self._generate_filenames(self._files, root = temp)
        stdin  = self._open_pipe(kwords, "IN_STDIN" , "rb")
        stdout = self._open_pipe(kwords, "OUT_STDOUT", "wb")
        stderr = self._open_pipe(kwords, "OUT_STDERR", "wb")

        cwd = None
        if self._set_cwd:
            cwd, kwords = temp, self._generate_filenames(self._files, root = None)

        command = [(field % kwords) for field in self._command]
        self._proc = subprocess.Popen(command,
                                      stdin  = stdin,
                                      stdout = stdout,
                                      stderr = stderr,
                                      cwd    = cwd,
                                      preexec_fn = os.setsid,
                                      close_fds = True)

        # Allow subprocesses to be killed in case of a SIGTERM
        _add_to_killlist(self._proc)

    def ready(self):
        """Returns true if the command has been run to completion,
        regardless of wether or not an error occured."""
        return (self._proc and self._proc.poll() is not None)

    def join(self):
        """Similar to Popen.wait(), but returns the value wrapped in a list,
        and ensures that any opened handles are closed. Must be called before
        calling commit."""
        try:
            return_codes = [self._proc.wait()] if self._proc else [None]
        finally:
            # Close any implictly opened pipes
            for (mode, handle) in self._handles:
                if "w" in mode:
                    handle.flush()
                handle.close()
            self._handles = []

        return return_codes

    def wait(self):
        """Equivalent to Subproces.wait. This function should only
        be used in contexts where a AtomicCmd needs to be combined
        with Subprocesses, as it does not exist for AtomicSets."""
        return self.join()[0]

    def terminate(self):
        """Sends SIGTERM to process."""
        self._proc.terminate()
        self._proc = None

    @property
    def executables(self):
        """Returns a list of executables required for the AtomicCmd."""
        return self._file_sets["executable"]

    @property
    def requirements(self):
        """Returns a list of callable requirements that must be met
        before the AtomicCmd can be expected to be runable."""
        return self._file_sets["requirements"]

    @property
    def input_files(self):
        """Returns a list of input files that are required by the AtomicCmd."""
        return self._file_sets["input"]

    @property
    def output_files(self):
        """Checks that the expected output files have been generated."""
        return self._file_sets["output"]

    @property
    def auxiliary_files(self):
        """Checks that the expected output files have been generated."""
        return self._file_sets["auxiliary"]

    def commit(self, temp):
        if not self.ready():
            raise CmdError("Attempting to commit command before it has completed")
        elif self._handles:
            raise CmdError("Called 'commit' before calling 'join'")

        self._proc = None

        temp_files = self._generate_filenames(self._files, temp)
        for (key, filename) in temp_files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("OUT_") and not os.path.exists(filename):
                    raise CmdError("Command did not create expected output file: " + filename)

        for (key, filename) in temp_files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("OUT_"):
                    fileutils.move_file(filename, self._files[key])
                elif key.startswith("TEMP_OUT_") and os.path.exists(filename):
                    os.remove(filename)

    @property
    def stdout(self):
        """Returns the 'stdout' value for the current process. If no
        such process has been started, the return value is None."""
        if self._proc:
            return self._proc.stdout

    def __str__(self):
        def describe_pipe(template, pipe):
            if isinstance(pipe, types.StringTypes):
                return template % pipe
            elif isinstance(pipe, AtomicCmd):
                return template % "[AtomicCmd]"
            elif (pipe == AtomicCmd.PIPE):
                return template % "[PIPE]"
            else:
                return ""

        kwords = self._generate_filenames(self._files, "${TEMP}")
        command = " ".join([(field % kwords) for field in self._command])

        stdin = self._files.get("IN_STDIN")
        if stdin:
            command += describe_pipe(" < %s", stdin)

        stdout = self._files.get("OUT_STDOUT")
        stderr = self._files.get("OUT_STDERR")
        if (stdout != stderr):
            if stdout:
                command += describe_pipe(" > %s", stdout)
            if stderr:
                command += describe_pipe(" 2> %s", stderr)
        elif stdout:
            command += describe_pipe(" &> %s", stdout)

        return "<%s>" % command

    @classmethod
    def _process_arguments(cls, proc_id, command, kwargs):
        cls._validate_pipes(kwargs)

        files = {}
        for (key, value) in kwargs.iteritems():
            if cls._validate_argument(key, value):
                files[key] = value

        executable = os.path.basename(command[0])
        for pipe in ("STDOUT", "STDERR"):
            if not (kwargs.get("OUT_" + pipe) or kwargs.get("TEMP_OUT_" + pipe)):
                filename = "pipe_%s_%i.%s" % (executable, proc_id, pipe.lower())
                files["TEMP_OUT_" + pipe] = filename

        output_files = collections.defaultdict(list)
        for (key, filename) in kwargs.iteritems():
            if key.startswith("TEMP_OUT_") or key.startswith("OUT_"):
                if isinstance(filename, types.StringTypes):
                    output_files[os.path.basename(filename)].append(key)

        for (filename, keys) in output_files.iteritems():
            if len(keys) > 1:
                raise CmdError("Same output filename (%s) is specified for multiple keys: %s" \
                                   % (filename, ", ".join(keys)))

        return files

    @classmethod
    def _validate_pipes(cls, kwargs):
        """Checks that no single pipe is specified multiple times, e.i. being specified
        both for a temporary and a final (outside the temp dir) file. For example,
        either IN_STDIN or TEMP_IN_STDIN must be specified, but not both."""
        if any((kwargs.get(pipe) and kwargs.get("TEMP_" + pipe)) for pipe in _PIPES):
            raise CmdError, "Pipes must be specified at most once (w/wo TEMP_)."

    @classmethod
    def _validate_argument(cls, key, value):
        if not any(key.startswith(prefix) for prefix in _PREFIXES):
            raise CmdError("Command contains invalid argument (wrong prefix): '%s' -> '%s'" \
                           % (cls.__name__, key))
        elif (value == cls.PIPE) and key not in ("OUT_STDOUT", "TEMP_OUT_STDOUT"):
            raise ValueError, "PIPE is only allow for *_STDOUT, not " + key
        elif key.startswith("CHECK_") and not isinstance(value, collections.Callable):
            raise ValueError("Values CHECKS_ must be callables, not %s" % repr(value))
        elif (key == "IN_STDIN") and not isinstance(value, types.StringTypes + (AtomicCmd,)):
            raise ValueError("STDIN for '%s' is not a string or AtomicCmd: %s" \
                             % (cls.__name__, value))

        # Values are allowed to be none, but such entries are skipped
        # This simplifies the use of keyword parameters with no default values.
        return bool(value)

    def _open_pipe(self, kwords, pipe, mode):
        filename = kwords.get(pipe, kwords.get("TEMP_" + pipe))
        if filename in (None, self.PIPE):
            return filename
        elif isinstance(filename, AtomicCmd):
            return filename.stdout

        pipe = open(filename, mode)
        self._handles.append((mode, pipe))

        return pipe

    @classmethod
    def _generate_filenames(cls, files, root):
        root = os.path.abspath(root) if root else ""

        filenames = {"TEMP_DIR" : root}
        for (key, filename) in files.iteritems():
            if isinstance(filename, types.StringTypes):
                if key.startswith("TEMP_") or key.startswith("OUT_"):
                    filename = os.path.join(root, os.path.basename(filename))
                elif not root and (key.startswith("IN_") or key.startswith("AUX_")):
                    filename = os.path.abspath(filename)
            filenames[key] = filename

        return filenames

    @classmethod
    def _build_files_map(cls, command, files):
        key_map   = {"IN" : "input",  "OUT": "output",  "EXEC" : "executable",
                     "AUX" : "auxiliary", "CHECK" : "requirements"}
        file_sets = dict((key, set()) for key in key_map.itervalues())

        file_sets["executable"].add(command[0])
        for (key, filename) in files.iteritems():
            if not isinstance(filename, types.StringTypes) and not key.startswith("CHECK_"):
                continue
            elif not key.startswith("TEMP_"):
                key = key_map[key.split("_", 1)[0]]
                file_sets[key].add(filename)

        return dict(zip(file_sets.keys(), map(frozenset, file_sets.values())))

## The following ensures proper cleanup of child processes, for example in the case
## where multiprocessing.Pool.terminate() is called. The current implementation will
## leak weakref objects, but given the low number of commands called over the lifetime
## of a pipeline, this is considered acceptable for now. FIXME
_PROCS = []

def _cleanup_children(signum, _frame):
    for proc_ref in _PROCS:
        proc = proc_ref()
        if proc:
            os.killpg(proc.pid, signal.SIGTERM)
    sys.exit(-signum)

def _add_to_killlist(proc):
    if not _PROCS:
        signal.signal(signal.SIGTERM, _cleanup_children)

    _PROCS.append(weakref.ref(proc))
<|endoftext|>