<commit_before>#***************************************************************************
#*   (c) sliptonic (shopinthewoods@gmail.com) 2014                        *
#*                                                                         *
#*   This file is part of the FreeCAD CAx development system.              *
#*                                                                         *
#*   This program is free software; you can redistribute it and/or modify  *
#*   it under the terms of the GNU Lesser General Public License (LGPL)    *
#*   as published by the Free Software Foundation; either version 2 of     *
#*   the License, or (at your option) any later version.                   *
#*   for detail see the LICENCE text file.                                 *
#*                                                                         *
#*   FreeCAD 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 Lesser General Public License for more details.                   *
#*                                                                         *
#*   You should have received a copy of the GNU Library General Public     *
#*   License along with FreeCAD; if not, write to the Free Software        *
#*   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  *
#*   USA                                                                   *
#*                                                                         *
#***************************************************************************/

TOOLTIP='''
Generate g-code from a Path that is compatible with the grbl controller.

import grbl_post
grbl_post.export(object,"/path/to/file.ncc")
'''

import FreeCAD
import PathScripts.PostUtils as PostUtils
import argparse
import datetime
import shlex

now = datetime.datetime.now()

parser = argparse.ArgumentParser(prog='grbl', add_help=False)
parser.add_argument('--header', action='store_true', help='output headers (default)')
parser.add_argument('--no-header', action='store_true', help='suppress header output')
parser.add_argument('--comments', action='store_true', help='output comment (default)')
parser.add_argument('--no-comments', action='store_true', help='suppress comment output')
parser.add_argument('--line-numbers', action='store_true', help='prefix with line numbers')
parser.add_argument('--no-line-numbers', action='store_true', help='don\'t prefix with line numbers (default)')
parser.add_argument('--show-editor', action='store_true', help='pop up editor before writing output (default)')
parser.add_argument('--no-show-editor', action='store_true', help='don\'t pop up editor before writing output')
parser.add_argument('--precision', default='4', help='number of digits of precision, default=4')
parser.add_argument('--preamble', help='set commands to be issued before the first command, default="G17\nG90"')
parser.add_argument('--postamble', help='set commands to be issued after the last command, default="M05\nG17 G90\n; M2"')
parser.add_argument('--tool-change', action='store_true', help='insert M6 for tool changes')
parser.add_argument('--no-tool-change', action='store_true', help='suppress all tool chanage commands (default)')

TOOLTIP_ARGS=parser.format_help()

#These globals set common customization preferences
OUTPUT_COMMENTS = True
OUTPUT_HEADER = True
OUTPUT_LINE_NUMBERS = False
OUTPUT_TOOL_CHANGE = False
SHOW_EDITOR = True
MODAL = False #if true commands are suppressed if the same as previous line.
COMMAND_SPACE = " "
LINENR = 100 #line number starting value

#These globals will be reflected in the Machine configuration of the project
UNITS = "G21" #G21 for metric, G20 for us standard
MACHINE_NAME = "GRBL"
CORNER_MIN = {'x':0, 'y':0, 'z':0 }
CORNER_MAX = {'x':500, 'y':300, 'z':300 }
PRECISION = 4

RAPID_MOVES = ['G0', 'G00']

#Preamble text will appear at the beginning of the GCODE output file.
PREAMBLE = '''G17 G90
'''

#Postamble text will appear following the last operation.
POSTAMBLE = '''M5
G17 G90
; M2
'''

# These commands are ignored by commenting them out
SUPPRESS_COMMANDS = [ 'G98', 'G80' ]

#Pre operation text will be inserted before every operation
PRE_OPERATION = ''''''

#Post operation text will be inserted after every operation
POST_OPERATION = ''''''

#Tool Change commands will be inserted before a tool change
TOOL_CHANGE = ''''''

# to distinguish python built-in open function from the one declared below
if open.__module__ == '__builtin__':
    pythonopen = open

def processArguments(argstring):
    global OUTPUT_HEADER
    global OUTPUT_COMMENTS
    global OUTPUT_LINE_NUMBERS
    global OUTPUT_TOOL_CHANGE
    global SHOW_EDITOR
    global PRECISION
    global PREAMBLE
    global POSTAMBLE

    try:
        args = parser.parse_args(shlex.split(argstring))
        if args.no_header:
            OUTPUT_HEADER = False
        if args.header:
            OUTPUT_HEADER = True
        if args.no_comments:
            OUTPUT_COMMENTS = False
        if args.comments:
            OUTPUT_COMMENTS = True
        if args.no_line_numbers:
            OUTPUT_LINE_NUMBERS = False
        if args.line_numbers:
            OUTPUT_LINE_NUMBERS = True
        if args.no_show_editor:
            SHOW_EDITOR = False
        if args.show_editor:
            SHOW_EDITOR = True
        print("Show editor = %d" % SHOW_EDITOR)
        PRECISION = args.precision
        if args.preamble is not None:
            PREAMBLE = args.preamble
        if args.postamble is not None:
            POSTAMBLE = args.postamble
        if args.tool_change:
            OUTPUT_TOOL_CHANGE = True
        if args.no_tool_change:
            OUTPUT_TOOL_CHANGE = False
    except:
        return False

    return True

def export(objectslist,filename,argstring):
    if not processArguments(argstring):
        return None

    global UNITS

    for obj in objectslist:
        if not hasattr(obj,"Path"):
            print("the object " + obj.Name + " is not a path. Please select only path and Compounds.")
            return

    print("postprocessing...")
    gcode = ""

    #Find the machine.
    #The user my have overridden post processor defaults in the GUI.  Make sure we're using the current values in the Machine Def.
    myMachine = None
    for pathobj in objectslist:
        if hasattr(pathobj,"Group"): #We have a compound or project.
            for p in pathobj.Group:
                if p.Name == "Machine":
                    myMachine = p
    if myMachine is None:
        print("No machine found in this project")
    else:
        if myMachine.MachineUnits == "Metric":
           UNITS = "G21"
        else:
           UNITS = "G20"

    # write header
    if OUTPUT_HEADER:
        gcode += linenumber() + "(Exported by FreeCAD)\n"
        gcode += linenumber() + "(Post Processor: " + __name__ +")\n"
        gcode += linenumber() + "(Output Time:"+str(now)+")\n"

    #Write the preamble
    if OUTPUT_COMMENTS: gcode += linenumber() + "(begin preamble)\n"
    for line in PREAMBLE.splitlines(True):
        gcode += linenumber() + line
    gcode += linenumber() + UNITS + "\n"

    for obj in objectslist:

        #do the pre_op
        if OUTPUT_COMMENTS: gcode += linenumber() + "(begin operation: " + obj.Label + ")\n"
        for line in PRE_OPERATION.splitlines(True):
            gcode += linenumber() + line

        gcode += parse(obj)

        #do the post_op
        if OUTPUT_COMMENTS: gcode += linenumber() + "(finish operation: " + obj.Label + ")\n"
        for line in POST_OPERATION.splitlines(True):
            gcode += linenumber() + line

    #do the post_amble

    if OUTPUT_COMMENTS: gcode += "(begin postamble)\n"
    for line in POSTAMBLE.splitlines(True):
        gcode += linenumber() + line

    if FreeCAD.GuiUp and SHOW_EDITOR:
        dia = PostUtils.GCodeEditorDialog()
        dia.editor.setText(gcode)
        result = dia.exec_()
        if result:
            final = dia.editor.toPlainText()
        else:
            final = gcode
    else:
        final = gcode

    print("done postprocessing.")

    gfile = pythonopen(filename,"wb")
    gfile.write(gcode)
    gfile.close()

def linenumber():
    global LINENR
    if OUTPUT_LINE_NUMBERS == True:
        LINENR += 10
        return "N" + str(LINENR) + " "
    return ""

def parse(pathobj):
    out = ""
    lastcommand = None
    precision_string = '.' + str(PRECISION) +'f'

    #params = ['X','Y','Z','A','B','I','J','K','F','S'] #This list control the order of parameters
    params = ['X','Y','Z','A','B','I','J','F','S','T','Q','R','L'] #linuxcnc doesn't want K properties on XY plane  Arcs need work.

    if hasattr(pathobj,"Group"): #We have a compound or project.
        if OUTPUT_COMMENTS: out += linenumber() + "(compound: " + pathobj.Label + ")\n"
        for p in pathobj.Group:
            out += parse(p)
        return out
    else: #parsing simple path

        if not hasattr(pathobj,"Path"): #groups might contain non-path things like stock.
            return out

        if OUTPUT_COMMENTS: out += linenumber() + "(Path: " + pathobj.Label + ")\n"

        for c in pathobj.Path.Commands:
            outstring = []
            command = c.Name
            outstring.append(command)
            # if modal: only print the command if it is not the same as the last one
            if MODAL == True:
                if command == lastcommand:
                    outstring.pop(0)

            # Now add the remaining parameters in order
            for param in params:
                if param in c.Parameters:
                    if param == 'F':
                        if command not in RAPID_MOVES:
                            outstring.append(param + format(c.Parameters['F'] * 60, '.2f'))
                    elif param == 'T':
                        outstring.append(param + str(int(c.Parameters['T'])))
                    else:
                        outstring.append(param + format(c.Parameters[param], precision_string))

            # store the latest command
            lastcommand = command

            # Check for Tool Change:
            if command == 'M6':
                if OUTPUT_COMMENTS:
                    out += linenumber() + "(begin toolchange)\n"
                if not OUTPUT_TOOL_CHANGE:
                    outstring.insert(0, ";")
                else:
                    for line in TOOL_CHANGE.splitlines(True):
                        out += linenumber() + line

            if command == "message":
                if OUTPUT_COMMENTS == False:
                    out = []
                else:
                    outstring.pop(0) #remove the command

            if command in SUPPRESS_COMMANDS:
                outstring.insert(0, ";")

            #prepend a line number and append a newline
            if len(outstring) >= 1:
                if OUTPUT_LINE_NUMBERS:
                    outstring.insert(0,(linenumber()))

                #append the line to the final output
                for w in outstring:
                    out += w + COMMAND_SPACE
                out = out.strip() + "\n"

        return out

print(__name__ + " gcode postprocessor loaded.")

<commit_msg>Allow user to suppress the initial tool change.<commit_after>#***************************************************************************
#*   (c) sliptonic (shopinthewoods@gmail.com) 2014                        *
#*                                                                         *
#*   This file is part of the FreeCAD CAx development system.              *
#*                                                                         *
#*   This program is free software; you can redistribute it and/or modify  *
#*   it under the terms of the GNU Lesser General Public License (LGPL)    *
#*   as published by the Free Software Foundation; either version 2 of     *
#*   the License, or (at your option) any later version.                   *
#*   for detail see the LICENCE text file.                                 *
#*                                                                         *
#*   FreeCAD 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 Lesser General Public License for more details.                   *
#*                                                                         *
#*   You should have received a copy of the GNU Library General Public     *
#*   License along with FreeCAD; if not, write to the Free Software        *
#*   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  *
#*   USA                                                                   *
#*                                                                         *
#***************************************************************************/

TOOLTIP='''
Generate g-code from a Path that is compatible with the grbl controller.

import grbl_post
grbl_post.export(object,"/path/to/file.ncc")
'''

import FreeCAD
import PathScripts.PostUtils as PostUtils
import argparse
import datetime
import shlex
import traceback

now = datetime.datetime.now()

parser = argparse.ArgumentParser(prog='grbl', add_help=False)
parser.add_argument('--header', action='store_true', help='output headers (default)')
parser.add_argument('--no-header', action='store_true', help='suppress header output')
parser.add_argument('--comments', action='store_true', help='output comment (default)')
parser.add_argument('--no-comments', action='store_true', help='suppress comment output')
parser.add_argument('--line-numbers', action='store_true', help='prefix with line numbers')
parser.add_argument('--no-line-numbers', action='store_true', help='don\'t prefix with line numbers (default)')
parser.add_argument('--show-editor', action='store_true', help='pop up editor before writing output (default)')
parser.add_argument('--no-show-editor', action='store_true', help='don\'t pop up editor before writing output')
parser.add_argument('--precision', default='4', help='number of digits of precision, default=4')
parser.add_argument('--preamble', help='set commands to be issued before the first command, default="G17\nG90"')
parser.add_argument('--postamble', help='set commands to be issued after the last command, default="M05\nG17 G90\n; M2"')
parser.add_argument('--tool-change', help='0 ... suppress all tool change commands\n1 ... insert M6 for all tool changes\n2 ... insert M6 for all tool changes except the initial tool')

TOOLTIP_ARGS=parser.format_help()

#These globals set common customization preferences
OUTPUT_COMMENTS = True
OUTPUT_HEADER = True
OUTPUT_LINE_NUMBERS = False
OUTPUT_TOOL_CHANGE = False
SHOW_EDITOR = True
MODAL = False #if true commands are suppressed if the same as previous line.
COMMAND_SPACE = " "
LINENR = 100 #line number starting value

#These globals will be reflected in the Machine configuration of the project
UNITS = "G21" #G21 for metric, G20 for us standard
MACHINE_NAME = "GRBL"
CORNER_MIN = {'x':0, 'y':0, 'z':0 }
CORNER_MAX = {'x':500, 'y':300, 'z':300 }
PRECISION = 4

RAPID_MOVES = ['G0', 'G00']

#Preamble text will appear at the beginning of the GCODE output file.
PREAMBLE = '''G17 G90
'''

#Postamble text will appear following the last operation.
POSTAMBLE = '''M5
G17 G90
; M2
'''

# These commands are ignored by commenting them out
SUPPRESS_COMMANDS = [ 'G98', 'G80' ]

#Pre operation text will be inserted before every operation
PRE_OPERATION = ''''''

#Post operation text will be inserted after every operation
POST_OPERATION = ''''''

#Tool Change commands will be inserted before a tool change
TOOL_CHANGE = ''''''
SUPPRESS_TOOL_CHANGE=0

# to distinguish python built-in open function from the one declared below
if open.__module__ == '__builtin__':
    pythonopen = open

def processArguments(argstring):
    global OUTPUT_HEADER
    global OUTPUT_COMMENTS
    global OUTPUT_LINE_NUMBERS
    global OUTPUT_TOOL_CHANGE
    global SHOW_EDITOR
    global PRECISION
    global PREAMBLE
    global POSTAMBLE
    global SUPPRESS_TOOL_CHANGE

    try:
        args = parser.parse_args(shlex.split(argstring))
        if args.no_header:
            OUTPUT_HEADER = False
        if args.header:
            OUTPUT_HEADER = True
        if args.no_comments:
            OUTPUT_COMMENTS = False
        if args.comments:
            OUTPUT_COMMENTS = True
        if args.no_line_numbers:
            OUTPUT_LINE_NUMBERS = False
        if args.line_numbers:
            OUTPUT_LINE_NUMBERS = True
        if args.no_show_editor:
            SHOW_EDITOR = False
        if args.show_editor:
            SHOW_EDITOR = True
        print("Show editor = %d" % SHOW_EDITOR)
        PRECISION = args.precision
        if not args.preamble is None:
            PREAMBLE = args.preamble
        if not args.postamble is None:
            POSTAMBLE = args.postamble
        if not args.tool_change is None:
            OUTPUT_TOOL_CHANGE = int(args.tool_change) > 0
            SUPPRESS_TOOL_CHANGE = min(1, int(args.tool_change) - 1)
    except Exception as e:
        traceback.print_exc(e)
        return False

    return True

def export(objectslist,filename,argstring):
    if not processArguments(argstring):
        return None

    global UNITS

    for obj in objectslist:
        if not hasattr(obj,"Path"):
            print("the object " + obj.Name + " is not a path. Please select only path and Compounds.")
            return

    print("postprocessing...")
    gcode = ""

    #Find the machine.
    #The user my have overridden post processor defaults in the GUI.  Make sure we're using the current values in the Machine Def.
    myMachine = None
    for pathobj in objectslist:
        if hasattr(pathobj,"Group"): #We have a compound or project.
            for p in pathobj.Group:
                if p.Name == "Machine":
                    myMachine = p
    if myMachine is None:
        print("No machine found in this project")
    else:
        if myMachine.MachineUnits == "Metric":
           UNITS = "G21"
        else:
           UNITS = "G20"

    # write header
    if OUTPUT_HEADER:
        gcode += linenumber() + "(Exported by FreeCAD)\n"
        gcode += linenumber() + "(Post Processor: " + __name__ +")\n"
        gcode += linenumber() + "(Output Time:"+str(now)+")\n"

    #Write the preamble
    if OUTPUT_COMMENTS: gcode += linenumber() + "(begin preamble)\n"
    for line in PREAMBLE.splitlines(True):
        gcode += linenumber() + line
    gcode += linenumber() + UNITS + "\n"

    for obj in objectslist:

        #do the pre_op
        if OUTPUT_COMMENTS: gcode += linenumber() + "(begin operation: " + obj.Label + ")\n"
        for line in PRE_OPERATION.splitlines(True):
            gcode += linenumber() + line

        gcode += parse(obj)

        #do the post_op
        if OUTPUT_COMMENTS: gcode += linenumber() + "(finish operation: " + obj.Label + ")\n"
        for line in POST_OPERATION.splitlines(True):
            gcode += linenumber() + line

    #do the post_amble

    if OUTPUT_COMMENTS: gcode += "(begin postamble)\n"
    for line in POSTAMBLE.splitlines(True):
        gcode += linenumber() + line

    if FreeCAD.GuiUp and SHOW_EDITOR:
        dia = PostUtils.GCodeEditorDialog()
        dia.editor.setText(gcode)
        result = dia.exec_()
        if result:
            final = dia.editor.toPlainText()
        else:
            final = gcode
    else:
        final = gcode

    print("done postprocessing.")

    gfile = pythonopen(filename,"wb")
    gfile.write(gcode)
    gfile.close()

def linenumber():
    global LINENR
    if OUTPUT_LINE_NUMBERS == True:
        LINENR += 10
        return "N" + str(LINENR) + " "
    return ""

def parse(pathobj):
    out = ""
    lastcommand = None
    precision_string = '.' + str(PRECISION) +'f'
    global SUPPRESS_TOOL_CHANGE

    #params = ['X','Y','Z','A','B','I','J','K','F','S'] #This list control the order of parameters
    params = ['X','Y','Z','A','B','I','J','F','S','T','Q','R','L'] #linuxcnc doesn't want K properties on XY plane  Arcs need work.

    if hasattr(pathobj,"Group"): #We have a compound or project.
        if OUTPUT_COMMENTS: out += linenumber() + "(compound: " + pathobj.Label + ")\n"
        for p in pathobj.Group:
            out += parse(p)
        return out
    else: #parsing simple path

        if not hasattr(pathobj,"Path"): #groups might contain non-path things like stock.
            return out

        if OUTPUT_COMMENTS: out += linenumber() + "(Path: " + pathobj.Label + ")\n"

        for c in pathobj.Path.Commands:
            outstring = []
            command = c.Name
            outstring.append(command)
            # if modal: only print the command if it is not the same as the last one
            if MODAL == True:
                if command == lastcommand:
                    outstring.pop(0)

            # Now add the remaining parameters in order
            for param in params:
                if param in c.Parameters:
                    if param == 'F':
                        if command not in RAPID_MOVES:
                            outstring.append(param + format(c.Parameters['F'] * 60, '.2f'))
                    elif param == 'T':
                        outstring.append(param + str(int(c.Parameters['T'])))
                    else:
                        outstring.append(param + format(c.Parameters[param], precision_string))

            # store the latest command
            lastcommand = command

            # Check for Tool Change:
            if command == 'M6':
                if OUTPUT_COMMENTS:
                    out += linenumber() + "(begin toolchange)\n"
                if not OUTPUT_TOOL_CHANGE or SUPPRESS_TOOL_CHANGE > 0:
                    outstring.insert(0, ";")
                    SUPPRESS_TOOL_CHANGE = SUPPRESS_TOOL_CHANGE - 1
                else:
                    for line in TOOL_CHANGE.splitlines(True):
                        out += linenumber() + line

            if command == "message":
                if OUTPUT_COMMENTS == False:
                    out = []
                else:
                    outstring.pop(0) #remove the command

            if command in SUPPRESS_COMMANDS:
                outstring.insert(0, ";")

            #prepend a line number and append a newline
            if len(outstring) >= 1:
                if OUTPUT_LINE_NUMBERS:
                    outstring.insert(0,(linenumber()))

                #append the line to the final output
                for w in outstring:
                    out += w + COMMAND_SPACE
                out = out.strip() + "\n"

        return out

print(__name__ + " gcode postprocessor loaded.")

<|endoftext|>


<commit_before>"""Tests for all FEP-based estimators in ``alchemlyb``.

"""
import pytest

import numpy as np
import pandas as pd

from alchemlyb.parsing import gmx, amber, namd
from alchemlyb.estimators import MBAR
from alchemlyb.estimators import BAR
import alchemtest.gmx
import alchemtest.amber
import alchemtest.namd

def gmx_benzene_coul_u_nk():
    dataset = alchemtest.gmx.load_benzene()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['Coulomb']])

    return u_nk

def gmx_benzene_vdw_u_nk():
    dataset = alchemtest.gmx.load_benzene()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['VDW']])

    return u_nk

def gmx_expanded_ensemble_case_1():
    dataset = alchemtest.gmx.load_expanded_ensemble_case_1()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_expanded_ensemble_case_2():
    dataset = alchemtest.gmx.load_expanded_ensemble_case_2()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_expanded_ensemble_case_3():
    dataset = alchemtest.gmx.load_expanded_ensemble_case_3()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_water_particle_with_total_energy():
    dataset = alchemtest.gmx.load_water_particle_with_total_energy()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_water_particle_with_potential_energy():
    dataset = alchemtest.gmx.load_water_particle_with_potential_energy()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_water_particle_without_energy():
    dataset = alchemtest.gmx.load_water_particle_without_energy()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def amber_bace_example_complex_vdw():
    dataset = alchemtest.amber.load_bace_example()

    u_nk = pd.concat([amber.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['complex']['vdw']])
    return u_nk

def namd_tyr2ala():
    dataset = alchemtest.namd.load_tyr2ala()
    u_nk1 = namd.extract_u_nk(dataset['data']['forward'][0], T=300)
    u_nk2 = namd.extract_u_nk(dataset['data']['backward'][0], T=300)

    # combine dataframes of fwd and rev directions
    u_nk1.replace(0, np.nan, inplace=True)
    u_nk1[u_nk1.isnull()] = u_nk2
    u_nk1.replace(np.nan, 0, inplace=True)
    u_nk = u_nk1.sort_index(level=u_nk1.index.names[1:])

    return u_nk

class FEPestimatorMixin:
    """Mixin for all FEP Estimator test classes.

    """

    def compare_delta_f(self, X_delta_f):
        est = self.cls().fit(X_delta_f[0])
        delta_f, d_delta_f = self.get_delta_f(est)

        assert X_delta_f[1] == pytest.approx(delta_f, rel=1e-3)
        assert X_delta_f[2] == pytest.approx(d_delta_f, rel=1e-3)

    def get_delta_f(self, est):
        return est.delta_f_.iloc[0, -1], est.d_delta_f_.iloc[0, -1]

class TestMBAR(FEPestimatorMixin):
    """Tests for MBAR.

    """
    cls = MBAR

    @pytest.mark.parametrize('X_delta_f', [
        (gmx_benzene_coul_u_nk(), 3.041, 0.02088),
        (gmx_benzene_vdw_u_nk(), -3.007, 0.04519),
        (gmx_expanded_ensemble_case_1(), 75.923, 0.14124),
        (gmx_expanded_ensemble_case_2(), 75.915, 0.14372),
        (gmx_expanded_ensemble_case_3(), 76.173, 0.11345),
        (gmx_water_particle_with_total_energy(), -11.680, 0.083655),
        (gmx_water_particle_with_potential_energy(), -11.675, 0.083589),
        (gmx_water_particle_without_energy(), -11.654, 0.083415),
        (amber_bace_example_complex_vdw(), 2.40200, 0.0618453),
    ])
    def test_mbar(self, X_delta_f):
        self.compare_delta_f(X_delta_f)

class TestBAR(FEPestimatorMixin):
    """Tests for BAR.

    """
    cls = BAR

    T = 300
    kT_NAMD = namd.k_b * T

    @pytest.mark.parametrize('X_delta_f', [
        (gmx_benzene_coul_u_nk(), 3.044, 0.01640),
        (gmx_benzene_vdw_u_nk(), -3.033, 0.03438),
        (gmx_expanded_ensemble_case_1(), 75.993, 0.11056),
        (gmx_expanded_ensemble_case_2(), 76.009, 0.11220),
        (gmx_expanded_ensemble_case_3(), 76.219, 0.08886),
        (gmx_water_particle_with_total_energy(), -11.675, 0.065055),
        (gmx_water_particle_with_potential_energy(), -11.724, 0.064964),
        (gmx_water_particle_without_energy(), -11.660, 0.064914),
        (amber_bace_example_complex_vdw(), 2.37846, 0.050899),
        (namd_tyr2ala(), 6.031269829/kT_NAMD, 0.069813058/kT_NAMD),
    ])
    def test_bar(self, X_delta_f):
        self.compare_delta_f(X_delta_f)

    def get_delta_f(self, est):
        ee = 0.0

        for i in range(len(est.d_delta_f_) - 1):
            ee += est.d_delta_f_.values[i][i+1]**2
        return est.delta_f_.iloc[0, -1], ee**0.5
<commit_msg>use parametrized fixtures in test_fep<commit_after>"""Tests for all FEP-based estimators in ``alchemlyb``.

"""
import pytest

import numpy as np
import pandas as pd

from alchemlyb.parsing import gmx, amber, namd
from alchemlyb.estimators import MBAR
from alchemlyb.estimators import BAR
import alchemtest.gmx
import alchemtest.amber
import alchemtest.namd

def gmx_benzene_coul_u_nk():
    dataset = alchemtest.gmx.load_benzene()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['Coulomb']])

    return u_nk

def gmx_benzene_vdw_u_nk():
    dataset = alchemtest.gmx.load_benzene()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['VDW']])

    return u_nk

def gmx_expanded_ensemble_case_1():
    dataset = alchemtest.gmx.load_expanded_ensemble_case_1()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_expanded_ensemble_case_2():
    dataset = alchemtest.gmx.load_expanded_ensemble_case_2()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_expanded_ensemble_case_3():
    dataset = alchemtest.gmx.load_expanded_ensemble_case_3()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_water_particle_with_total_energy():
    dataset = alchemtest.gmx.load_water_particle_with_total_energy()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_water_particle_with_potential_energy():
    dataset = alchemtest.gmx.load_water_particle_with_potential_energy()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def gmx_water_particle_without_energy():
    dataset = alchemtest.gmx.load_water_particle_without_energy()

    u_nk = pd.concat([gmx.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['AllStates']])

    return u_nk

def amber_bace_example_complex_vdw():
    dataset = alchemtest.amber.load_bace_example()

    u_nk = pd.concat([amber.extract_u_nk(filename, T=300)
                      for filename in dataset['data']['complex']['vdw']])
    return u_nk

def namd_tyr2ala():
    dataset = alchemtest.namd.load_tyr2ala()
    u_nk1 = namd.extract_u_nk(dataset['data']['forward'][0], T=300)
    u_nk2 = namd.extract_u_nk(dataset['data']['backward'][0], T=300)

    # combine dataframes of fwd and rev directions
    u_nk1.replace(0, np.nan, inplace=True)
    u_nk1[u_nk1.isnull()] = u_nk2
    u_nk1.replace(np.nan, 0, inplace=True)
    u_nk = u_nk1.sort_index(level=u_nk1.index.names[1:])

    return u_nk

class FEPestimatorMixin:
    """Mixin for all FEP Estimator test classes.

    """

    def compare_delta_f(self, X_delta_f):
        est = self.cls().fit(X_delta_f[0])
        delta_f, d_delta_f = self.get_delta_f(est)

        assert X_delta_f[1] == pytest.approx(delta_f, rel=1e-3)
        assert X_delta_f[2] == pytest.approx(d_delta_f, rel=1e-3)

    def get_delta_f(self, est):
        return est.delta_f_.iloc[0, -1], est.d_delta_f_.iloc[0, -1]

class TestMBAR(FEPestimatorMixin):
    """Tests for MBAR.

    """
    cls = MBAR

    @pytest.fixture(scope="class",
                    params=[(gmx_benzene_coul_u_nk, 3.041, 0.02088),
                            (gmx_benzene_vdw_u_nk, -3.007, 0.04519),
                            (gmx_expanded_ensemble_case_1, 75.923, 0.14124),
                            (gmx_expanded_ensemble_case_2, 75.915, 0.14372),
                            (gmx_expanded_ensemble_case_3, 76.173, 0.11345),
                            (gmx_water_particle_with_total_energy, -11.680, 0.083655),
                            (gmx_water_particle_with_potential_energy, -11.675, 0.083589),
                            (gmx_water_particle_without_energy, -11.654, 0.083415),
                            (amber_bace_example_complex_vdw, 2.40200, 0.0618453),
                    ])
    def X_delta_f(self, request):
        get_unk, E, dE = request.param
        return get_unk(), E, dE

    def test_mbar(self, X_delta_f):
        self.compare_delta_f(X_delta_f)

class TestBAR(FEPestimatorMixin):
    """Tests for BAR.

    """
    cls = BAR

    T = 300
    kT_NAMD = namd.k_b * T

    @pytest.fixture(scope="class",
                    params = [(gmx_benzene_coul_u_nk, 3.044, 0.01640),
                              (gmx_benzene_vdw_u_nk, -3.033, 0.03438),
                              (gmx_expanded_ensemble_case_1, 75.993, 0.11056),
                              (gmx_expanded_ensemble_case_2, 76.009, 0.11220),
                              (gmx_expanded_ensemble_case_3, 76.219, 0.08886),
                              (gmx_water_particle_with_total_energy, -11.675, 0.065055),
                              (gmx_water_particle_with_potential_energy, -11.724, 0.064964),
                              (gmx_water_particle_without_energy, -11.660, 0.064914),
                              (amber_bace_example_complex_vdw, 2.37846, 0.050899),
                              (namd_tyr2ala, 6.031269829/kT_NAMD, 0.069813058/kT_NAMD),
                    ])
    def X_delta_f(self, request):
        get_unk, E, dE = request.param
        return get_unk(), E, dE

    def test_bar(self, X_delta_f):
        self.compare_delta_f(X_delta_f)

    def get_delta_f(self, est):
        ee = 0.0

        for i in range(len(est.d_delta_f_) - 1):
            ee += est.d_delta_f_.values[i][i+1]**2
        return est.delta_f_.iloc[0, -1], ee**0.5
<|endoftext|>


<commit_before># --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------

from __future__ import print_function
import argparse

from azure.cli.core.commands import ExtensionCommandSource
from azure.cli.core.commands.constants import SURVEY_PROMPT, SURVEY_PROMPT_COLOR

from knack.help import (HelpFile as KnackHelpFile, CommandHelpFile as KnackCommandHelpFile,
                        GroupHelpFile as KnackGroupHelpFile, ArgumentGroupRegistry as KnackArgumentGroupRegistry,
                        HelpExample as KnackHelpExample, HelpParameter as KnackHelpParameter,
                        _print_indent, CLIHelp, HelpAuthoringException)

from knack.log import get_logger
from knack.util import CLIError

logger = get_logger(__name__)

PRIVACY_STATEMENT = """
Welcome to Azure CLI!
---------------------
Use `az -h` to see available commands or go to https://aka.ms/cli.

Telemetry
---------
The Azure CLI collects usage data in order to improve your experience.
The data is anonymous and does not include commandline argument values.
The data is collected by Microsoft.

You can change your telemetry settings with `az configure`.
"""

WELCOME_MESSAGE = r"""
     /\
    /  \    _____   _ _  ___ _
   / /\ \  |_  / | | | \'__/ _\
  / ____ \  / /| |_| | | |  __/
 /_/    \_\/___|\__,_|_|  \___|

Welcome to the cool new Azure CLI!

Use `az --version` to display the current version.
Here are the base commands:
"""

# PrintMixin class to decouple printing functionality from AZCLIHelp class.
# Most of these methods override print methods in CLIHelp
class CLIPrintMixin(CLIHelp):
    def _print_header(self, cli_name, help_file):
        super(CLIPrintMixin, self)._print_header(cli_name, help_file)

        links = help_file.links
        if links:
            link_text = "{} and {}".format(", ".join([link["url"] for link in links[0:-1]]),
                                           links[-1]["url"]) if len(links) > 1 else links[0]["url"]
            link_text = "For more information, see: {}\n".format(link_text)
            _print_indent(link_text, 2, width=self.textwrap_width)

    def _print_detailed_help(self, cli_name, help_file):
        CLIPrintMixin._print_extensions_msg(help_file)
        super(CLIPrintMixin, self)._print_detailed_help(cli_name, help_file)
        self._print_az_find_message(help_file.command, self.cli_ctx.enable_color)

    @staticmethod
    def _get_choices_defaults_sources_str(p):
        choice_str = u'  Allowed values: {}.'.format(', '.join(sorted([str(x) for x in p.choices]))) \
            if p.choices else ''
        default_str = u'  Default: {}.'.format(p.default) if p.default and p.default != argparse.SUPPRESS else ''
        value_sources_str = CLIPrintMixin._process_value_sources(p) if p.value_sources else ''
        return u'{}{}{}'.format(choice_str, default_str, value_sources_str)

    @staticmethod
    def _print_examples(help_file):
        indent = 0
        _print_indent('Examples', indent)
        for e in help_file.examples:
            indent = 1
            _print_indent(u'{0}'.format(e.short_summary), indent)
            indent = 2
            if e.long_summary:
                _print_indent(u'{0}'.format(e.long_summary), indent)
            _print_indent(u'{0}'.format(e.command), indent)
            print('')

    @staticmethod
    def _print_az_find_message(command, enable_color):
        from colorama import Style
        indent = 0
        message = 'For more specific examples, use: az find "az {}"'.format(command)
        if enable_color:
            message = Style.BRIGHT + message + Style.RESET_ALL
        _print_indent(message + '\n', indent)

    @staticmethod
    def _process_value_sources(p):
        commands, strings, urls = [], [], []

        for item in p.value_sources:
            if "string" in item:
                strings.append(item["string"])
            elif "link" in item and "command" in item["link"]:
                commands.append(item["link"]["command"])
            elif "link" in item and "url" in item["link"]:
                urls.append(item["link"]["url"])

        command_str = u'  Values from: {}.'.format(", ".join(commands)) if commands else ''
        string_str = u'  {}'.format(", ".join(strings)) if strings else ''
        string_str = string_str + "." if string_str and not string_str.endswith(".") else string_str
        urls_str = u'  For more info, go to: {}.'.format(", ".join(urls)) if urls else ''
        return u'{}{}{}'.format(command_str, string_str, urls_str)

    @staticmethod
    def _print_extensions_msg(help_file):
        if help_file.type != 'command':
            return
        if isinstance(help_file.command_source, ExtensionCommandSource):
            logger.warning(help_file.command_source.get_command_warn_msg())
            # If experimental is true, it overrides preview
            if help_file.command_source.experimental:
                logger.warning(help_file.command_source.get_experimental_warn_msg())
            elif help_file.command_source.preview:
                logger.warning(help_file.command_source.get_preview_warn_msg())

class AzCliHelp(CLIPrintMixin, CLIHelp):

    def __init__(self, cli_ctx):
        super(AzCliHelp, self).__init__(cli_ctx,
                                        privacy_statement=PRIVACY_STATEMENT,
                                        welcome_message=WELCOME_MESSAGE,
                                        command_help_cls=CliCommandHelpFile,
                                        group_help_cls=CliGroupHelpFile,
                                        help_cls=CliHelpFile)
        from knack.help import HelpObject

        # TODO: This workaround is used to avoid a bizarre bug in Python 2.7. It
        # essentially reassigns Knack's HelpObject._normalize_text implementation
        # with an identical implemenation in Az. For whatever reason, this fixes
        # the bug in Python 2.7.
        @staticmethod
        def new_normalize_text(s):
            if not s or len(s) < 2:
                return s or ''
            s = s.strip()
            initial_upper = s[0].upper() + s[1:]
            trailing_period = '' if s[-1] in '.!?' else '.'
            return initial_upper + trailing_period

        HelpObject._normalize_text = new_normalize_text  # pylint: disable=protected-access

        self._register_help_loaders()
        self._name_to_content = {}

    def show_help(self, cli_name, nouns, parser, is_group):
        self.update_loaders_with_help_file_contents(nouns)

        delimiters = ' '.join(nouns)
        help_file = self.command_help_cls(self, delimiters, parser) if not is_group \
            else self.group_help_cls(self, delimiters, parser)
        help_file.load(parser)
        if not nouns:
            help_file.command = ''
        else:
            AzCliHelp.update_examples(help_file)
        self._print_detailed_help(cli_name, help_file)

        print(SURVEY_PROMPT_COLOR if self.cli_ctx.enable_color else SURVEY_PROMPT)

    def _register_help_loaders(self):
        import azure.cli.core._help_loaders as help_loaders
        import inspect

        def is_loader_cls(cls):
            return inspect.isclass(cls) and cls.__name__ != 'BaseHelpLoader'and issubclass(cls, help_loaders.BaseHelpLoader)  # pylint: disable=line-too-long

        versioned_loaders = {}
        for cls_name, loader_cls in inspect.getmembers(help_loaders, is_loader_cls):
            loader = loader_cls(self)
            versioned_loaders[cls_name] = loader

        if len(versioned_loaders) != len({ldr.version for ldr in versioned_loaders.values()}):
            ldrs_str = " ".join("{}-version:{}".format(cls_name, ldr.version) for cls_name, ldr in versioned_loaders.items())  # pylint: disable=line-too-long
            raise CLIError("Two loaders have the same version. Loaders:\n\t{}".format(ldrs_str))

        self.versioned_loaders = versioned_loaders

    def update_loaders_with_help_file_contents(self, nouns):
        loader_file_names_dict = {}
        file_name_set = set()
        for ldr_cls_name, loader in self.versioned_loaders.items():
            new_file_names = loader.get_noun_help_file_names(nouns) or []
            loader_file_names_dict[ldr_cls_name] = new_file_names
            file_name_set.update(new_file_names)

        for file_name in file_name_set:
            if file_name not in self._name_to_content:
                with open(file_name, 'r') as f:
                    self._name_to_content[file_name] = f.read()

        for ldr_cls_name, file_names in loader_file_names_dict.items():
            file_contents = {}
            for name in file_names:
                file_contents[name] = self._name_to_content[name]
            self.versioned_loaders[ldr_cls_name].update_file_contents(file_contents)

    # This method is meant to be a hook that can be overridden by an extension or module.
    @staticmethod
    def update_examples(help_file):
        pass

class CliHelpFile(KnackHelpFile):

    def __init__(self, help_ctx, delimiters):
        # Each help file (for a command or group) has a version denoting the source of its data.
        super(CliHelpFile, self).__init__(help_ctx, delimiters)
        self.links = []

    def _should_include_example(self, ex):
        supported_profiles = ex.get('supported-profiles')
        unsupported_profiles = ex.get('unsupported-profiles')

        if all((supported_profiles, unsupported_profiles)):
            raise HelpAuthoringException("An example cannot have both supported-profiles and unsupported-profiles.")

        if supported_profiles:
            supported_profiles = [profile.strip() for profile in supported_profiles.split(',')]
            return self.help_ctx.cli_ctx.cloud.profile in supported_profiles

        if unsupported_profiles:
            unsupported_profiles = [profile.strip() for profile in unsupported_profiles.split(',')]
            return self.help_ctx.cli_ctx.cloud.profile not in unsupported_profiles

        return True

    # Needs to override base implementation to exclude unsupported examples.
    def _load_from_data(self, data):
        if not data:
            return

        if isinstance(data, str):
            self.long_summary = data
            return

        if 'type' in data:
            self.type = data['type']

        if 'short-summary' in data:
            self.short_summary = data['short-summary']

        self.long_summary = data.get('long-summary')

        if 'examples' in data:
            self.examples = []
            for d in data['examples']:
                if self._should_include_example(d):
                    self.examples.append(HelpExample(**d))

    def load(self, options):
        ordered_loaders = sorted(self.help_ctx.versioned_loaders.values(), key=lambda ldr: ldr.version)
        for loader in ordered_loaders:
            loader.versioned_load(self, options)

class CliGroupHelpFile(KnackGroupHelpFile, CliHelpFile):

    def load(self, options):
        # forces class to use this load method even if KnackGroupHelpFile overrides CliHelpFile's method.
        CliHelpFile.load(self, options)

class CliCommandHelpFile(KnackCommandHelpFile, CliHelpFile):

    def __init__(self, help_ctx, delimiters, parser):
        super(CliCommandHelpFile, self).__init__(help_ctx, delimiters, parser)
        self.type = 'command'
        self.command_source = getattr(parser, 'command_source', None)

        self.parameters = []

        for action in [a for a in parser._actions if a.help != argparse.SUPPRESS]:  # pylint: disable=protected-access
            if action.option_strings:
                self._add_parameter_help(action)
            else:
                # use metavar for positional parameters
                param_kwargs = {
                    'name_source': [action.metavar or action.dest],
                    'deprecate_info': getattr(action, 'deprecate_info', None),
                    'preview_info': getattr(action, 'preview_info', None),
                    'experimental_info': getattr(action, 'experimental_info', None),
                    'description': action.help,
                    'choices': action.choices,
                    'required': False,
                    'default': None,
                    'group_name': 'Positional'
                }
                self.parameters.append(HelpParameter(**param_kwargs))

        help_param = next(p for p in self.parameters if p.name == '--help -h')
        help_param.group_name = 'Global Arguments'

        # update parameter type so we can use overriden update_from_data method to update value sources.
        for param in self.parameters:
            param.__class__ = HelpParameter

    def _load_from_data(self, data):
        super(CliCommandHelpFile, self)._load_from_data(data)

        if isinstance(data, str) or not self.parameters or not data.get('parameters'):
            return

        loaded_params = []
        loaded_param = {}
        for param in self.parameters:
            loaded_param = next((n for n in data['parameters'] if n['name'] == param.name), None)
            if loaded_param:
                param.update_from_data(loaded_param)
            loaded_params.append(param)

        self.parameters = loaded_params

    def load(self, options):
        # forces class to use this load method even if KnackCommandHelpFile overrides CliHelpFile's method.
        CliHelpFile.load(self, options)

class ArgumentGroupRegistry(KnackArgumentGroupRegistry):  # pylint: disable=too-few-public-methods

    def __init__(self, group_list):

        super(ArgumentGroupRegistry, self).__init__(group_list)
        self.priorities = {
            None: 0,
            'Resource Id Arguments': 1,
            'Generic Update Arguments': 998,
            'Global Arguments': 1000,
        }
        priority = 2
        # any groups not already in the static dictionary should be prioritized alphabetically
        other_groups = [g for g in sorted(list(set(group_list))) if g not in self.priorities]
        for group in other_groups:
            self.priorities[group] = priority
            priority += 1

class HelpExample(KnackHelpExample):  # pylint: disable=too-few-public-methods

    def __init__(self, **_data):
        # Old attributes
        _data['name'] = _data.get('name', '')
        _data['text'] = _data.get('text', '')
        super(HelpExample, self).__init__(_data)

        self.name = _data.get('summary', '') if _data.get('summary', '') else self.name
        self.text = _data.get('command', '') if _data.get('command', '') else self.text

        self.long_summary = _data.get('description', '')
        self.supported_profiles = _data.get('supported-profiles', None)
        self.unsupported_profiles = _data.get('unsupported-profiles', None)

    # alias old params with new
    @property
    def short_summary(self):
        return self.name

    @short_summary.setter
    def short_summary(self, value):
        self.name = value

    @property
    def command(self):
        return self.text

    @command.setter
    def command(self, value):
        self.text = value

class HelpParameter(KnackHelpParameter):  # pylint: disable=too-many-instance-attributes

    def __init__(self, **kwargs):
        super(HelpParameter, self).__init__(**kwargs)

    def update_from_data(self, data):
        super(HelpParameter, self).update_from_data(data)
        # original help.py value_sources are strings, update command strings to value-source dict
        if self.value_sources:
            self.value_sources = [str_or_dict if isinstance(str_or_dict, dict) else {"link": {"command": str_or_dict}}
                                  for str_or_dict in self.value_sources]
<commit_msg>{Core} Disable extension preview/experimental warning (#12656)<commit_after># --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------

from __future__ import print_function
import argparse

from azure.cli.core.commands import ExtensionCommandSource
from azure.cli.core.commands.constants import SURVEY_PROMPT, SURVEY_PROMPT_COLOR

from knack.help import (HelpFile as KnackHelpFile, CommandHelpFile as KnackCommandHelpFile,
                        GroupHelpFile as KnackGroupHelpFile, ArgumentGroupRegistry as KnackArgumentGroupRegistry,
                        HelpExample as KnackHelpExample, HelpParameter as KnackHelpParameter,
                        _print_indent, CLIHelp, HelpAuthoringException)

from knack.log import get_logger
from knack.util import CLIError

logger = get_logger(__name__)

PRIVACY_STATEMENT = """
Welcome to Azure CLI!
---------------------
Use `az -h` to see available commands or go to https://aka.ms/cli.

Telemetry
---------
The Azure CLI collects usage data in order to improve your experience.
The data is anonymous and does not include commandline argument values.
The data is collected by Microsoft.

You can change your telemetry settings with `az configure`.
"""

WELCOME_MESSAGE = r"""
     /\
    /  \    _____   _ _  ___ _
   / /\ \  |_  / | | | \'__/ _\
  / ____ \  / /| |_| | | |  __/
 /_/    \_\/___|\__,_|_|  \___|

Welcome to the cool new Azure CLI!

Use `az --version` to display the current version.
Here are the base commands:
"""

# PrintMixin class to decouple printing functionality from AZCLIHelp class.
# Most of these methods override print methods in CLIHelp
class CLIPrintMixin(CLIHelp):
    def _print_header(self, cli_name, help_file):
        super(CLIPrintMixin, self)._print_header(cli_name, help_file)

        links = help_file.links
        if links:
            link_text = "{} and {}".format(", ".join([link["url"] for link in links[0:-1]]),
                                           links[-1]["url"]) if len(links) > 1 else links[0]["url"]
            link_text = "For more information, see: {}\n".format(link_text)
            _print_indent(link_text, 2, width=self.textwrap_width)

    def _print_detailed_help(self, cli_name, help_file):
        CLIPrintMixin._print_extensions_msg(help_file)
        super(CLIPrintMixin, self)._print_detailed_help(cli_name, help_file)
        self._print_az_find_message(help_file.command, self.cli_ctx.enable_color)

    @staticmethod
    def _get_choices_defaults_sources_str(p):
        choice_str = u'  Allowed values: {}.'.format(', '.join(sorted([str(x) for x in p.choices]))) \
            if p.choices else ''
        default_str = u'  Default: {}.'.format(p.default) if p.default and p.default != argparse.SUPPRESS else ''
        value_sources_str = CLIPrintMixin._process_value_sources(p) if p.value_sources else ''
        return u'{}{}{}'.format(choice_str, default_str, value_sources_str)

    @staticmethod
    def _print_examples(help_file):
        indent = 0
        _print_indent('Examples', indent)
        for e in help_file.examples:
            indent = 1
            _print_indent(u'{0}'.format(e.short_summary), indent)
            indent = 2
            if e.long_summary:
                _print_indent(u'{0}'.format(e.long_summary), indent)
            _print_indent(u'{0}'.format(e.command), indent)
            print('')

    @staticmethod
    def _print_az_find_message(command, enable_color):
        from colorama import Style
        indent = 0
        message = 'For more specific examples, use: az find "az {}"'.format(command)
        if enable_color:
            message = Style.BRIGHT + message + Style.RESET_ALL
        _print_indent(message + '\n', indent)

    @staticmethod
    def _process_value_sources(p):
        commands, strings, urls = [], [], []

        for item in p.value_sources:
            if "string" in item:
                strings.append(item["string"])
            elif "link" in item and "command" in item["link"]:
                commands.append(item["link"]["command"])
            elif "link" in item and "url" in item["link"]:
                urls.append(item["link"]["url"])

        command_str = u'  Values from: {}.'.format(", ".join(commands)) if commands else ''
        string_str = u'  {}'.format(", ".join(strings)) if strings else ''
        string_str = string_str + "." if string_str and not string_str.endswith(".") else string_str
        urls_str = u'  For more info, go to: {}.'.format(", ".join(urls)) if urls else ''
        return u'{}{}{}'.format(command_str, string_str, urls_str)

    @staticmethod
    def _print_extensions_msg(help_file):
        if help_file.type != 'command':
            return
        if isinstance(help_file.command_source, ExtensionCommandSource):
            logger.warning(help_file.command_source.get_command_warn_msg())

            # Extension preview/experimental warning is disabled because it can be confusing when displayed together
            # with command or command group preview/experimental warning. See #12556

            # # If experimental is true, it overrides preview
            # if help_file.command_source.experimental:
            #     logger.warning(help_file.command_source.get_experimental_warn_msg())
            # elif help_file.command_source.preview:
            #     logger.warning(help_file.command_source.get_preview_warn_msg())

class AzCliHelp(CLIPrintMixin, CLIHelp):

    def __init__(self, cli_ctx):
        super(AzCliHelp, self).__init__(cli_ctx,
                                        privacy_statement=PRIVACY_STATEMENT,
                                        welcome_message=WELCOME_MESSAGE,
                                        command_help_cls=CliCommandHelpFile,
                                        group_help_cls=CliGroupHelpFile,
                                        help_cls=CliHelpFile)
        from knack.help import HelpObject

        # TODO: This workaround is used to avoid a bizarre bug in Python 2.7. It
        # essentially reassigns Knack's HelpObject._normalize_text implementation
        # with an identical implemenation in Az. For whatever reason, this fixes
        # the bug in Python 2.7.
        @staticmethod
        def new_normalize_text(s):
            if not s or len(s) < 2:
                return s or ''
            s = s.strip()
            initial_upper = s[0].upper() + s[1:]
            trailing_period = '' if s[-1] in '.!?' else '.'
            return initial_upper + trailing_period

        HelpObject._normalize_text = new_normalize_text  # pylint: disable=protected-access

        self._register_help_loaders()
        self._name_to_content = {}

    def show_help(self, cli_name, nouns, parser, is_group):
        self.update_loaders_with_help_file_contents(nouns)

        delimiters = ' '.join(nouns)
        help_file = self.command_help_cls(self, delimiters, parser) if not is_group \
            else self.group_help_cls(self, delimiters, parser)
        help_file.load(parser)
        if not nouns:
            help_file.command = ''
        else:
            AzCliHelp.update_examples(help_file)
        self._print_detailed_help(cli_name, help_file)

        print(SURVEY_PROMPT_COLOR if self.cli_ctx.enable_color else SURVEY_PROMPT)

    def _register_help_loaders(self):
        import azure.cli.core._help_loaders as help_loaders
        import inspect

        def is_loader_cls(cls):
            return inspect.isclass(cls) and cls.__name__ != 'BaseHelpLoader'and issubclass(cls, help_loaders.BaseHelpLoader)  # pylint: disable=line-too-long

        versioned_loaders = {}
        for cls_name, loader_cls in inspect.getmembers(help_loaders, is_loader_cls):
            loader = loader_cls(self)
            versioned_loaders[cls_name] = loader

        if len(versioned_loaders) != len({ldr.version for ldr in versioned_loaders.values()}):
            ldrs_str = " ".join("{}-version:{}".format(cls_name, ldr.version) for cls_name, ldr in versioned_loaders.items())  # pylint: disable=line-too-long
            raise CLIError("Two loaders have the same version. Loaders:\n\t{}".format(ldrs_str))

        self.versioned_loaders = versioned_loaders

    def update_loaders_with_help_file_contents(self, nouns):
        loader_file_names_dict = {}
        file_name_set = set()
        for ldr_cls_name, loader in self.versioned_loaders.items():
            new_file_names = loader.get_noun_help_file_names(nouns) or []
            loader_file_names_dict[ldr_cls_name] = new_file_names
            file_name_set.update(new_file_names)

        for file_name in file_name_set:
            if file_name not in self._name_to_content:
                with open(file_name, 'r') as f:
                    self._name_to_content[file_name] = f.read()

        for ldr_cls_name, file_names in loader_file_names_dict.items():
            file_contents = {}
            for name in file_names:
                file_contents[name] = self._name_to_content[name]
            self.versioned_loaders[ldr_cls_name].update_file_contents(file_contents)

    # This method is meant to be a hook that can be overridden by an extension or module.
    @staticmethod
    def update_examples(help_file):
        pass

class CliHelpFile(KnackHelpFile):

    def __init__(self, help_ctx, delimiters):
        # Each help file (for a command or group) has a version denoting the source of its data.
        super(CliHelpFile, self).__init__(help_ctx, delimiters)
        self.links = []

    def _should_include_example(self, ex):
        supported_profiles = ex.get('supported-profiles')
        unsupported_profiles = ex.get('unsupported-profiles')

        if all((supported_profiles, unsupported_profiles)):
            raise HelpAuthoringException("An example cannot have both supported-profiles and unsupported-profiles.")

        if supported_profiles:
            supported_profiles = [profile.strip() for profile in supported_profiles.split(',')]
            return self.help_ctx.cli_ctx.cloud.profile in supported_profiles

        if unsupported_profiles:
            unsupported_profiles = [profile.strip() for profile in unsupported_profiles.split(',')]
            return self.help_ctx.cli_ctx.cloud.profile not in unsupported_profiles

        return True

    # Needs to override base implementation to exclude unsupported examples.
    def _load_from_data(self, data):
        if not data:
            return

        if isinstance(data, str):
            self.long_summary = data
            return

        if 'type' in data:
            self.type = data['type']

        if 'short-summary' in data:
            self.short_summary = data['short-summary']

        self.long_summary = data.get('long-summary')

        if 'examples' in data:
            self.examples = []
            for d in data['examples']:
                if self._should_include_example(d):
                    self.examples.append(HelpExample(**d))

    def load(self, options):
        ordered_loaders = sorted(self.help_ctx.versioned_loaders.values(), key=lambda ldr: ldr.version)
        for loader in ordered_loaders:
            loader.versioned_load(self, options)

class CliGroupHelpFile(KnackGroupHelpFile, CliHelpFile):

    def load(self, options):
        # forces class to use this load method even if KnackGroupHelpFile overrides CliHelpFile's method.
        CliHelpFile.load(self, options)

class CliCommandHelpFile(KnackCommandHelpFile, CliHelpFile):

    def __init__(self, help_ctx, delimiters, parser):
        super(CliCommandHelpFile, self).__init__(help_ctx, delimiters, parser)
        self.type = 'command'
        self.command_source = getattr(parser, 'command_source', None)

        self.parameters = []

        for action in [a for a in parser._actions if a.help != argparse.SUPPRESS]:  # pylint: disable=protected-access
            if action.option_strings:
                self._add_parameter_help(action)
            else:
                # use metavar for positional parameters
                param_kwargs = {
                    'name_source': [action.metavar or action.dest],
                    'deprecate_info': getattr(action, 'deprecate_info', None),
                    'preview_info': getattr(action, 'preview_info', None),
                    'experimental_info': getattr(action, 'experimental_info', None),
                    'description': action.help,
                    'choices': action.choices,
                    'required': False,
                    'default': None,
                    'group_name': 'Positional'
                }
                self.parameters.append(HelpParameter(**param_kwargs))

        help_param = next(p for p in self.parameters if p.name == '--help -h')
        help_param.group_name = 'Global Arguments'

        # update parameter type so we can use overriden update_from_data method to update value sources.
        for param in self.parameters:
            param.__class__ = HelpParameter

    def _load_from_data(self, data):
        super(CliCommandHelpFile, self)._load_from_data(data)

        if isinstance(data, str) or not self.parameters or not data.get('parameters'):
            return

        loaded_params = []
        loaded_param = {}
        for param in self.parameters:
            loaded_param = next((n for n in data['parameters'] if n['name'] == param.name), None)
            if loaded_param:
                param.update_from_data(loaded_param)
            loaded_params.append(param)

        self.parameters = loaded_params

    def load(self, options):
        # forces class to use this load method even if KnackCommandHelpFile overrides CliHelpFile's method.
        CliHelpFile.load(self, options)

class ArgumentGroupRegistry(KnackArgumentGroupRegistry):  # pylint: disable=too-few-public-methods

    def __init__(self, group_list):

        super(ArgumentGroupRegistry, self).__init__(group_list)
        self.priorities = {
            None: 0,
            'Resource Id Arguments': 1,
            'Generic Update Arguments': 998,
            'Global Arguments': 1000,
        }
        priority = 2
        # any groups not already in the static dictionary should be prioritized alphabetically
        other_groups = [g for g in sorted(list(set(group_list))) if g not in self.priorities]
        for group in other_groups:
            self.priorities[group] = priority
            priority += 1

class HelpExample(KnackHelpExample):  # pylint: disable=too-few-public-methods

    def __init__(self, **_data):
        # Old attributes
        _data['name'] = _data.get('name', '')
        _data['text'] = _data.get('text', '')
        super(HelpExample, self).__init__(_data)

        self.name = _data.get('summary', '') if _data.get('summary', '') else self.name
        self.text = _data.get('command', '') if _data.get('command', '') else self.text

        self.long_summary = _data.get('description', '')
        self.supported_profiles = _data.get('supported-profiles', None)
        self.unsupported_profiles = _data.get('unsupported-profiles', None)

    # alias old params with new
    @property
    def short_summary(self):
        return self.name

    @short_summary.setter
    def short_summary(self, value):
        self.name = value

    @property
    def command(self):
        return self.text

    @command.setter
    def command(self, value):
        self.text = value

class HelpParameter(KnackHelpParameter):  # pylint: disable=too-many-instance-attributes

    def __init__(self, **kwargs):
        super(HelpParameter, self).__init__(**kwargs)

    def update_from_data(self, data):
        super(HelpParameter, self).update_from_data(data)
        # original help.py value_sources are strings, update command strings to value-source dict
        if self.value_sources:
            self.value_sources = [str_or_dict if isinstance(str_or_dict, dict) else {"link": {"command": str_or_dict}}
                                  for str_or_dict in self.value_sources]
<|endoftext|>


<commit_before>"""
This is a base-class for writing custom external lookups.

This is licensed under the Apache License Version 2.0
See https://www.apache.org/licenses/LICENSE-2.0.html

To use this, you will need to:

 1) A lookup command script that implements from CustomLookup
 2) Create a transforms.conf that declares your lookup command
     2.1) The path of your lookup command filename (e.g. echo_lookup.py) and the names of the fields that need to be provided
     2.2) The list of fields that the command will generate (usually will include the field passed into it)

See below for an example. Note that this command can be run with:

* | head 1 | lookup echo host

default/transforms.conf:
--------------------------------

[echo]
external_cmd = echo_lookup.py host
fields_list = host,echohost

bin/echo_lookup.py:
--------------------------------

import logging
from custom_lookup import CustomLookup

class EchoLookup(CustomLookup):
    def __init__(self):
        CustomLookup.__init__(self, ['echohost'], 'echo_lookup_command', logging.INFO)

    def do_lookup(self, host):
        return {'echohost' : host}

EchoLookup.main()
"""

import csv
import sys
import logging
from logging import handlers

from splunk.appserver.mrsparkle.lib.util import make_splunkhome_path

class CustomLookup(object):

    def __init__(self, fieldnames=None, logger_name='custom_lookup_command', log_level=logging.INFO):
        """
        Constructs an instance of the lookup command.

        Arguments:
        fieldnames -- A list of the field names that the command will output
        logger_name -- The logger name to append to the logger
        log_level -- The log level to use for the logger
        """

        # Check and save the logger name
        self._logger = None

        if logger_name is None or len(logger_name) == 0:
            raise Exception("Logger name cannot be empty")

        self.logger_name = logger_name
        self.log_level = log_level

        # Keep a list of the invalid fields so that we don't re-warn people about the same things
        self.invalid_fields = set()

        # Ensure that the list of the accepted fieldnames is valid
        if fieldnames is None or len(fieldnames) == 0:
            raise Exception("The value for fieldnames must include an array of at least one row")

        # Here is a list of the accepted fieldnames
        self.fieldnames = fieldnames

    @property
    def logger(self):
        """
        A property that returns the logger.
        """

        # Make a logger unless it already exists
        if self._logger is not None:
            return self._logger

        logger = logging.getLogger(self.logger_name)

         # Prevent the log messages from being duplicated in the python.log file:
        logger.propagate = False
        logger.setLevel(self.log_level)

        file_handler = handlers.RotatingFileHandler(make_splunkhome_path(['var', 'log', 'splunk', self.logger_name + '.log']), maxBytes=25000000, backupCount=5)
        formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
        file_handler.setFormatter(formatter)

        logger.addHandler(file_handler)

        self._logger = logger
        return self._logger

    @logger.setter
    def logger(self, logger):
        self._logger = logger

    def do_lookup(self, *args, **kwargs):
        """
        This is the function that performs the lookup. It must be sub-classed.
        """
        raise Exception("do_lookup needs to be implemented")

    def extend_fieldnames(self, fieldnames):
        """
        Extends the list of fieldnames with those that the lookup will create.
        """

        fieldnames.extend(self.fieldnames)
        return set(fieldnames)

    def add_result(self, result_dict, output_dict, fieldnames, only_overwrite_empty=False):
        """
        Merge the output into the result dictionary but don't merge in fields that aren't listed in
        fieldnames.
        """

        # Keep a count of the invalid fields so that we can detect if we found new ones
        start_invalid_fields_len = len(self.invalid_fields)

        # Add each field to the output
        for field, value in output_dict.items():

            # Don't overwrite fields with an existing value since this could cause Splunk to not
            # be able to match the results with the original values.
            # see https://lukemurphey.net/issues/2348
            if only_overwrite_empty and result_dict.get(field, 'DOESNT EXIST YET') not in ['', None]:
                self.logger.debug("Skipping blank field %r ", field)
                pass

            # Make sure that field is in the list of field names
            elif field in fieldnames:

                self.logger.debug("Adding field %s with value %r", field, value)

                # If the output is a list, then include a comma deliminated list
                if isinstance(value, (list, tuple)) and not isinstance(value, basestring):
                    result_dict[field] = ", ".join(value)

                # The entry is a flat value, just include it
                else:
                    result_dict[field] = value

            # Detected an unexpected value, log it
            else:
                self.invalid_fields.add(field)

        # Output a warning if there are fields we didn't expect
        if len(self.invalid_fields) > start_invalid_fields_len:
            self.logger.warn("Discovered field that are not in the fields list: %r", self.invalid_fields)

    @classmethod
    def main(cls):
        lookup_command = None

        try:
            lookup_command = cls()
            lookup_command.execute()

        except Exception as e:

            # This logs general exceptions that would have been unhandled otherwise (such as coding
            # errors)
            if lookup_command is not None and lookup_command.logger is not None:
                lookup_command.logger.exception("Unhandled exception was caught, this may be due to a defect in the script")
            else:
                raise e

    def execute(self):
        """
        Execute the lookup command based on the values from standard input and output.
        """

        # This contains a list of the fields to perform an operation on (like "host")
        args = sys.argv[1:]

        # Initialize the input for the results
        infile = sys.stdin
        r = csv.DictReader(infile)
        # fieldnames = self.extend_fieldnames(r.fieldnames)
        # We need to use the original header because Splunk gets confused if you add fields to the
        # header that it didn't already have. It then fails to match up the input with the output
        # fields.
        # see https://lukemurphey.net/issues/2348
        fieldnames = r.fieldnames

        # Initialize the output for the results
        outfile = sys.stdout
        w = csv.DictWriter(outfile, fieldnames=fieldnames)
        w.writeheader()

        # Process each result
        for result in r:

            # Make up the arguments for the lookup call
            keyword_arguments = {}

            for parameter_name in args:
                keyword_arguments[parameter_name] = result.get(parameter_name, None)

            # Perform the lookup
            output = self.do_lookup(**keyword_arguments)

            # Put the output in the result
            if output:
                self.add_result(result, output, fieldnames, False)

            # Write out the result
            w.writerow(result)<commit_msg>Making the lookup base class support multi-threading<commit_after>"""
This is a base-class for writing custom external lookups.

This is licensed under the Apache License Version 2.0
See https://www.apache.org/licenses/LICENSE-2.0.html

To use this, you will need to:

 1) A lookup command script that implements from CustomLookup
 2) Create a transforms.conf that declares your lookup command
     2.1) The path of your lookup command filename (e.g. echo_lookup.py) and the names of the fields that need to be provided
     2.2) The list of fields that the command will generate (usually will include the field passed into it)

See below for an example. Note that this command can be run with:

* | head 1 | lookup echo host

default/transforms.conf:
--------------------------------

[echo]
external_cmd = echo_lookup.py host
fields_list = host,echohost

bin/echo_lookup.py:
--------------------------------

import logging
from custom_lookup import CustomLookup

class EchoLookup(CustomLookup):
    def __init__(self):
        CustomLookup.__init__(self, ['echohost'], 'echo_lookup_command', logging.INFO)

    def do_lookup(self, host):
        return {'echohost' : host}

EchoLookup.main()
"""

import csv
import sys
import time
import logging
from logging import handlers
import threading

from splunk.appserver.mrsparkle.lib.util import make_splunkhome_path

class CustomLookup(object):

    def __init__(self, fieldnames=None, logger_name='custom_lookup_command', log_level=logging.INFO, thread_limit=25):
        """
        Constructs an instance of the lookup command.

        Arguments:
        fieldnames -- A list of the field names that the command will output
        logger_name -- The logger name to append to the logger
        log_level -- The log level to use for the logger
        """

        # Check and save the logger name
        self._logger = None

        if logger_name is None or len(logger_name) == 0:
            raise Exception("Logger name cannot be empty")

        self.logger_name = logger_name
        self.log_level = log_level

        # Keep a list of the invalid fields so that we don't re-warn people about the same things
        self.invalid_fields = set()

        # Ensure that the list of the accepted fieldnames is valid
        if fieldnames is None or len(fieldnames) == 0:
            raise Exception("The value for fieldnames must include an array of at least one row")

        # Here is a list of the accepted fieldnames
        self.fieldnames = fieldnames

        # Use threading if requested
        self.thread_limit = thread_limit
        self.threads = []
        self.lock = threading.RLock()
        self.semaphore = threading.BoundedSemaphore(self.thread_limit)

    @property
    def logger(self):
        """
        A property that returns the logger.
        """

        # Make a logger unless it already exists
        if self._logger is not None:
            return self._logger

        logger = logging.getLogger(self.logger_name)

         # Prevent the log messages from being duplicated in the python.log file:
        logger.propagate = False
        logger.setLevel(self.log_level)

        file_handler = handlers.RotatingFileHandler(make_splunkhome_path(['var', 'log', 'splunk', self.logger_name + '.log']), maxBytes=25000000, backupCount=5)
        formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
        file_handler.setFormatter(formatter)

        logger.addHandler(file_handler)

        self._logger = logger
        return self._logger

    @logger.setter
    def logger(self, logger):
        self._logger = logger

    def do_lookup(self, *args, **kwargs):
        """
        This is the function that performs the lookup. It must be sub-classed.
        """
        raise Exception("do_lookup needs to be implemented")

    def extend_fieldnames(self, fieldnames):
        """
        Extends the list of fieldnames with those that the lookup will create.
        """

        fieldnames.extend(self.fieldnames)
        return set(fieldnames)

    def add_result(self, result_dict, output_dict, fieldnames, only_overwrite_empty=False):
        """
        Merge the output into the result dictionary but don't merge in fields that aren't listed in
        fieldnames.
        """

        # Keep a count of the invalid fields so that we can detect if we found new ones
        start_invalid_fields_len = len(self.invalid_fields)

        # Add each field to the output
        for field, value in output_dict.items():

            # Don't overwrite fields with an existing value since this could cause Splunk to not
            # be able to match the results with the original values.
            # see https://lukemurphey.net/issues/2348
            if only_overwrite_empty and result_dict.get(field, 'DOESNT EXIST YET') not in ['', None]:
                self.logger.debug("Skipping blank field %r ", field)
                pass

            # Make sure that field is in the list of field names
            elif field in fieldnames:

                self.logger.debug("Adding field %s with value %r", field, value)

                # If the output is a list, then include a comma deliminated list
                if isinstance(value, (list, tuple)) and not isinstance(value, basestring):
                    result_dict[field] = ", ".join(value)

                # The entry is a flat value, just include it
                else:
                    result_dict[field] = value

            # Detected an unexpected value, log it
            else:
                self.invalid_fields.add(field)

        # Output a warning if there are fields we didn't expect
        if len(self.invalid_fields) > start_invalid_fields_len:
            self.logger.warn("Discovered field that are not in the fields list: %r", self.invalid_fields)

    @classmethod
    def main(cls):
        lookup_command = None

        try:
            lookup_command = cls()
            lookup_command.execute()

            lookup_command.logger.info("Starting lookup execution")

        except Exception as e:

            # This logs general exceptions that would have been unhandled otherwise (such as coding
            # errors)
            if lookup_command is not None and lookup_command.logger is not None:
                lookup_command.logger.exception("Unhandled exception was caught, this may be due to a defect in the script")
            else:
                raise e

    def execute_lookup(self, input_row, output_csv, fieldnames):
        """
        Execute the lookup command on the given row and put the result in the output
        """

        # Try to acquire the semaphore which enforces the thread limit
        with self.semaphore:

            # This contains a list of the fields to perform an operation on (like "host")
            args = sys.argv[1:]

            # Make up the arguments for the lookup call
            keyword_arguments = {}

            for parameter_name in args:
                keyword_arguments[parameter_name] = input_row.get(parameter_name, None)

            # Perform the lookup
            output = self.do_lookup(**keyword_arguments)

            # Put the output in the result
            if output:
                self.add_result(input_row, output, fieldnames, False)

            # Write out the result
            with self.lock:
                output_csv.writerow(input_row)

    def execute(self):
        """
        Execute the lookup command based on the values from standard input and output.
        """

        # Record the start time
        start_time = time.time()

        # Initialize the input for the results
        infile = sys.stdin
        r = csv.DictReader(infile)

        # We need to use the original header because Splunk gets confused if you add fields to the
        # header that it didn't already have. It then fails to match up the input with the output
        # fields.
        # see https://lukemurphey.net/issues/2348
        fieldnames = r.fieldnames

        # Initialize the output for the results
        outfile = sys.stdout
        w = csv.DictWriter(outfile, fieldnames=fieldnames)
        w.writeheader()

        # Process each result
        for result in r:

            if self.thread_limit <= 1:
                self.execute_lookup(result, w, fieldnames)
            else:
                def do_lookup():
                    self.execute_lookup(result, w, fieldnames)

                # Start a thread
                new_thread = threading.Thread(name='ping_lookup', target=do_lookup)
                self.threads.append(new_thread)
                new_thread.start()

        # Wait for the threads to complete before moving on
        for thread in self.threads:
            thread.join()

        self.logger.info("Lookup completed, runtime=%ss", round(time.time() - start_time, 2))
<|endoftext|>


<commit_before><commit_msg>iniciando a aparecer dados corretos<commit_after><|endoftext|>