<commit_before>#!/usr/bin/python

import pywikibot
import re
from datetime import date
import datetime
import logging
import sys

########
# Changing this to 1 makes your changes live on the report page, do not set to
# live mode unless you have been approved for bot usage. Do not merge commits
# where this is not default to 0
########
live = 0
########
# Style: 0 = with date sectionsl 1 = without
########
style = 0
########
# Version Number
########
version = '0.4.6'
########

'''
Copyright (c) 2016 Wugpodes

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.
'''

def monthConvert(name):
    '''
    Takes in either the name of the month or the number of the month and returns
    the opposite. An input of str(July) would return int(7) while an input of
    int(6) would return str(June).
    Takes:   int OR string
    Returns: string OR int
    '''
    if type(name) is str:
        if name == "January": return 1
        elif name == "February": return 2
        elif name == "March": return 3
        elif name == "April": return 4
        elif name == "May": return 5
        elif name == "June": return 6
        elif name == "July": return 7
        elif name == "August": return 8
        elif name == "September": return 9
        elif name == "October": return 10
        elif name == "November": return 11
        elif name == "December": return 12
        else: raise ValueError
    elif type(name) is int:
        if name == 1:return('January')
        elif name == 2:return('February')
        elif name == 3:return('March')
        elif name == 4:return('April')
        elif name == 5:return('May')
        elif name == 6:return('June')
        elif name == 7:return('July')
        elif name == 8:return('August')
        elif name == 9:return('September')
        elif name == 10:return('October')
        elif name == 11: return('November')
        elif name == 12: return('December')
        else: raise ValueError

def checkPage(title):
    '''
    Takes the title of a DYK nom page and loads the wikitext. Looks for a DYK
    tick that shows it was approved or approved (AGF). If so it adds it to
    various lists so that it can be added to the approved page.
    Returns None, implicitly
    '''
    global dates
    global nonDate
    global entries
    global problem
    title = title.lstrip('{').rstrip('}')
    approved = computeNomStatus(title)
    if approved == -1:
        entries.pop()
    elif approved == 1:
        dates[-1][1].append('{{'+title+'}}')
        nonDate.append('{{'+title+'}}')
        entries.pop()
    else:
        pass

def computeNomStatus(link,status=0):
    link = link.lstrip('{').rstrip('}')
    if 'Template:' not in link:
        link = 'Template:'+link
    page = pywikibot.Page(site,link)
    if 'Please do not modify this page.' in page.text:
        status=-1
    else:
        for line in page.text.split('\n'):
            if '[[File:Symbol confirmed.svg|16px]]' in line \
            or '[[File:Symbol voting keep.svg|16px]]' in line:
                status = 1
            elif '[[File:Symbol question.svg|16px]]' in line \
            or '[[File:Symbol possible vote.svg|16px]]' in line \
            or '[[File:Symbol delete vote.svg|16px]]' in line \
            or '[[File:Symbol redirect vote 4.svg|16px]]' in line:
                status = 0
    return(status)

def mergeNominations(item = ''):
    '''
    Takes an index for the list dates and appends the entries in that section
    to the proper section on the approved page. Returns None, implicitly
    '''
    global approvedPageDates
    global dates
    if item == '':
        pass
    else:
        for entry in approvedPageDates:
            if dates[item][0] in entry:
                entry[1]+=dates[item][1][1:]

def computeYear(oMonth):
    '''
    Gives correct year over one year change, ie, for 11 months. Sections older
    than 11 months will not yield the correct results
    '''
    now = datetime.datetime.now()
    cMonth = now.month
    cYear = now.year
    if oMonth <= cMonth:
        year = cYear
    else:
        year = cYear - 1
    return(year)

def startLogging(loglevel):
    numeric_level = getattr(logging, loglevel.upper(), None)
    if numeric_level != None:
        logging.basicConfig(
            level=numeric_level,
            filename='Testing.log',
            format='%(asctime)s - %(levelname)s - %(message)s',
            datefmt='%Y-%m-%d %I:%M:%S %p'
        )
    else:
        logging.basicConfig(
            level=logging.WARNING,
            filename='./log/DYKMoverBot.log',
            format='%(asctime)s - %(levelname)s - %(message)s',
            datefmt='%Y-%m-%d %I:%M:%S %p'
        )
        logging.warning('Invalid log level \'%s\'. ' \
                       +'Defaulting to WARNING' % loglevel)

def checkArgs(arg):
    arg = arg.split('=')
    if arg[0] == '--log' or arg[0] == '-l':
        startLogging(arg[1])
    else:
        raise ValueError('Unknown command line argument \'%s\'' % arg[0])

# Start log
#for i in range(1,len(sys.argv)):
#    checkArgs(sys.argv[i])

logging.info("### Starting new run ###")
logging.info("live is set to %s" % live)
logging.info("style is set to %s" % style)
logging.info("DYKMoverBot version %s" % version)

# Load the various pages
logging.info("Loading pages")
site = pywikibot.Site('en', 'wikipedia')
nomPage      = pywikibot.Page(site,'Template talk:Did you know')
#approvedPage = pywikibot.Page(site,'Template talk:Did you know/Approved')
approvedPage = pywikibot.Page(site,'User:Wugapodes/DYKTest/0')
approvedPage1= pywikibot.Page(site,'User:Wugapodes/DYKTest/1')

dateRegex = re.compile(r'on (.*?) (\d+)=')

DYKpage = nomPage.text.split('\n')
dates   = []
nonDate = []
entries = []
problem = []

# Get approved nominations on the DYK page
logging.info("Reading Nomination page")
for line in DYKpage:
    entries.append(line)
    if '==Articles' in line:
        matches=dateRegex.search(line)
        month = monthConvert(str(matches.group(1)))
        day = int(matches.group(2))
        dt = datetime.date(month=month,day=day,year=computeYear(month))
        section = line
        dates.append([dt,[section]])
    elif 'Did you know nominations/' in line and '<!--' not in line:
        if '}}{{' in line:
            logging.debug("Multiple nominations on one line.\n"\
                         +"Section: %s" % line)
            splitLine = line.split('}}{{')
            for title in splitLine:
                try:
                    checkPage(title)
                except Exception as e:
                    logging.error(
                        "Function checkPage did not complete successfully.\n"\
                        +str(e)+"\n"+title
                    )
                    logging.debug(line)
                    print(e)
                    problem.append([line,e])
                    continue
        else:
            line = line.split('}')[0]
            try:
                checkPage(line)
            except Exception as e:
                logging.error(
                    "Function checkPage did not complete successfully.\n"\
                    +str(e)+"\n"+line
                )
                print(e)
                problem.append([line,e])
                continue

# Get the text on the approved page and then update the proper sections
# with entries newly approved
logging.info("Reading Approved page")
approvedPageText = approvedPage.text.split('\n')
sectionName = ''
approvedPageDates = []
oldLine = ''
datesToRemove = []
if style != 1:
    logging.info("Starting style 0 parsing")
    for line in approvedPageText:
        if '==Articles' in line:
            mergeNominations(oldLine)
            matches=dateRegex.search(line)
            month = monthConvert(str(matches.group(1)))
            day = int(matches.group(2))
            dt = datetime.date(month=month,day=day,year=computeYear(month))
            sectionName = line
            for item in dates:
                if dt in item:
                    oldLine = dates.index(item)
                    datesToRemove.append(oldLine)
                    break
                else:
                    continue
            approvedPageDates.append([dt,[sectionName]])
        elif 'Did you know nominations/' in line and '<!--' not in line:
            if '}}{{' in line:
                splitLine = line.split('}}{{')
                for title in splitLine:
                    if computeNomStatus(title) >= 0:
                        try:
                            approvedPageDates[-1][1].append('{{'+title+'}}')
                        except Exception as e:
                            logging.warning("The approved page is empty?\n"+str(e))
            else:
                line = line.split('}')[0]
                if computeNomStatus(title) >= 0:
                    try:
                        approvedPageDates[-1][1].append('{{'+title+'}}')
                    except Exception as e:
                        logging.warning("The approved page is empty?\n"+str(e))
    dates2 = []
    for i in datesToRemove:
        dates2.append(dates[i])
    newSections = [x for x in dates if x not in dates2 and len(x[1]) > 1]
    approvedPageDates+=newSections
    toPrint=[]
    approvedPageDates.sort(key=lambda x: x[0])
    for entry in dates:
        if len(entry[1]) > 1:
            toPrint+=entry[1]

approvedPage1Text = approvedPage1.text.split('\n')
toRemoveFromNonDate = []
if style != 0:
    logging.info("Starting style 1 parsing.")
    logging.warning("Style 1 is no longer maintained. "\
        +"Bot may behave incorrectly or not at all.")
    for line in approvedPage1Text:
        if line in nonDate:
            toRemoveFromNonDate.append(line)
    nonDateTemp = [x for x in nonDate if x not in toRemoveFromNonDate]
    nonDate = nonDateTemp

# Create the page text to be output
logging.info("Creating output text")
approvedText2 = approvedPage1.text.split('\n')
passed = 0
approvedText = [
        [
        "{{/top}}\n",
        "=Nominations=\n",
        "==Approved nominations==\n",
        "<!-- This section will hold approved nominations, with the templates "\
        +"transcluded in the same manner as the regular nominations page. "\
        +"While the exact format of the section has not yet been decided-while"\
        +" it seems unlikely that it will be by date, it may be divided into "\
        +"other sections-it is likely that the oldest approvals will go at the"\
        +" top and the most recent ones at the bottom of each section. -->\n"
        ],
        [
        "{{/top}}\n",
        "=Nominations=\n",
        "==Approved nominations==\n",
        "<!-- This section will hold approved nominations, with the templates "\
        +"transcluded in the same manner as the regular nominations page. "\
        +"While the exact format of the section has not yet been decided-while"\
        +" it seems unlikely that it will be by date, it may be divided into "\
        +"other sections-it is likely that the oldest approvals will go at the"\
        +" top and the most recent ones at the bottom of each section. -->\n"
        ]
    ]

for line in approvedText2:
    if '==Approved nominations==' in line:
        passed=1
    elif '==Special occasion holding area==' in line:
        passed=2
        approvedText[1].append('\n'.join(nonDate))
        approvedText[1].append('\n'+line+'\n')
    elif passed > 0:
        approvedText[1].append(line+'\n')
approvedText[0].append('\n'.join(toPrint))
for line in approvedPage.text.split('\n'):
    if '==Special occasion holding area==' in line:
        passed = 1
        approvedText[0].append('\n'+line+'\n')
    elif passed == 1:
        approvedText[0].append(line+'\n')

print('Done')

approvedText = [
        [
        "{{/top}}\n",
        "=Nominations=\n",
        "==Approved nominations==\n",
        "<!-- This section will hold approved nominations, with the templates "\
        +"transcluded in the same manner as the regular nominations page. "\
        +"While the exact format of the section has not yet been decided-while"\
        +" it seems unlikely that it will be by date, it may be divided into "\
        +"other sections-it is likely that the oldest approvals will go at the"\
        +" top and the most recent ones at the bottom of each section. -->\n"
        ],
        [
        "{{/top}}\n",
        "=Nominations=\n",
        "==Approved nominations==\n",
        "<!-- This section will hold approved nominations, with the templates "\
        +"transcluded in the same manner as the regular nominations page. "\
        +"While the exact format of the section has not yet been decided-while"\
        +" it seems unlikely that it will be by date, it may be divided into "\
        +"other sections-it is likely that the oldest approvals will go at the"\
        +" top and the most recent ones at the bottom of each section. -->\n"
        ]
    ]

# Determine if the bot should write to a live page or the test page. Defaults to
#     test page. Value of -1 tests backlog update (not standard because the file
#     size is very big).
if type(live) is str:
    pass
    page = pywikibot.Page(site,'Template talk:Did you know/Approved')
    page.text=''.join(approvedText[style])
    page.save('moving '+str(len(nonDate))+'tentatively approved nominations '\
              +'from [[WP:DYKN]], WugBot v'+version)
    page = pywikibot.Page(site,'Template talk:Did you know')
    # This is where the rest of the live stuff will go once I figure out
    #   how the other bots handle empty sections
else:
    page = pywikibot.Page(site,'User:Wugapodes/DYKTest/0')
    page.text=''.join(approvedText[0])
    page.save('test of DYKMover, WugBot v'+version+' section style')
    page = pywikibot.Page(site,'User:Wugapodes/DYKTest/1')
    page.text=''.join(approvedText[1])
    page.save('test of DYKMover, WugBot v'+version+' nosection style')
    page = pywikibot.Page(site,'User:Wugapodes/DYKTest/0')

print('Done')
<commit_msg>fixed blank page bug, v0.4.7<commit_after>#!/usr/bin/python

import pywikibot
import re
from datetime import date
import datetime
import logging
import sys

########
# Changing this to 1 makes your changes live on the report page, do not set to
# live mode unless you have been approved for bot usage. Do not merge commits
# where this is not default to 0
########
live = 0
########
# Style: 0 = with date sectionsl 1 = without
########
style = 0
########
# Version Number
########
version = '0.4.7'
########

'''
Copyright (c) 2016 Wugpodes

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.
'''

def monthConvert(name):
    '''
    Takes in either the name of the month or the number of the month and returns
    the opposite. An input of str(July) would return int(7) while an input of
    int(6) would return str(June).
    Takes:   int OR string
    Returns: string OR int
    '''
    if type(name) is str:
        if name == "January": return 1
        elif name == "February": return 2
        elif name == "March": return 3
        elif name == "April": return 4
        elif name == "May": return 5
        elif name == "June": return 6
        elif name == "July": return 7
        elif name == "August": return 8
        elif name == "September": return 9
        elif name == "October": return 10
        elif name == "November": return 11
        elif name == "December": return 12
        else: raise ValueError
    elif type(name) is int:
        if name == 1:return('January')
        elif name == 2:return('February')
        elif name == 3:return('March')
        elif name == 4:return('April')
        elif name == 5:return('May')
        elif name == 6:return('June')
        elif name == 7:return('July')
        elif name == 8:return('August')
        elif name == 9:return('September')
        elif name == 10:return('October')
        elif name == 11: return('November')
        elif name == 12: return('December')
        else: raise ValueError

def checkPage(title):
    '''
    Takes the title of a DYK nom page and loads the wikitext. Looks for a DYK
    tick that shows it was approved or approved (AGF). If so it adds it to
    various lists so that it can be added to the approved page.
    Returns None, implicitly
    '''
    global dates
    global nonDate
    global entries
    global problem
    title = title.lstrip('{').rstrip('}')
    approved = computeNomStatus(title)
    if approved == -1:
        entries.pop()
    elif approved == 1:
        dates[-1][1].append('{{'+title+'}}')
        nonDate.append('{{'+title+'}}')
        entries.pop()
    else:
        pass

def computeNomStatus(link,status=0):
    link = link.lstrip('{').rstrip('}')
    if 'Template:' not in link:
        link = 'Template:'+link
    page = pywikibot.Page(site,link)
    if 'Please do not modify this page.' in page.text:
        status=-1
    else:
        for line in page.text.split('\n'):
            if '[[File:Symbol confirmed.svg|16px]]' in line \
            or '[[File:Symbol voting keep.svg|16px]]' in line:
                status = 1
            elif '[[File:Symbol question.svg|16px]]' in line \
            or '[[File:Symbol possible vote.svg|16px]]' in line \
            or '[[File:Symbol delete vote.svg|16px]]' in line \
            or '[[File:Symbol redirect vote 4.svg|16px]]' in line:
                status = 0
    return(status)

def mergeNominations(item = ''):
    '''
    Takes an index for the list dates and appends the entries in that section
    to the proper section on the approved page. Returns None, implicitly
    '''
    global approvedPageDates
    global dates
    if item == '':
        pass
    else:
        for entry in approvedPageDates:
            if dates[item][0] in entry:
                entry[1]+=dates[item][1][1:]

def computeYear(oMonth):
    '''
    Gives correct year over one year change, ie, for 11 months. Sections older
    than 11 months will not yield the correct results
    '''
    now = datetime.datetime.now()
    cMonth = now.month
    cYear = now.year
    if oMonth <= cMonth:
        year = cYear
    else:
        year = cYear - 1
    return(year)

def startLogging(loglevel):
    numeric_level = getattr(logging, loglevel.upper(), None)
    if numeric_level != None:
        logging.basicConfig(
            level=numeric_level,
            filename='Testing.log',
            format='%(asctime)s - %(levelname)s - %(message)s',
            datefmt='%Y-%m-%d %I:%M:%S %p'
        )
    else:
        logging.basicConfig(
            level=logging.WARNING,
            filename='./log/DYKMoverBot.log',
            format='%(asctime)s - %(levelname)s - %(message)s',
            datefmt='%Y-%m-%d %I:%M:%S %p'
        )
        logging.warning('Invalid log level \'%s\'. ' \
                       +'Defaulting to WARNING' % loglevel)

def checkArgs(arg):
    arg = arg.split('=')
    if arg[0] == '--log' or arg[0] == '-l':
        startLogging(arg[1])
    else:
        raise ValueError('Unknown command line argument \'%s\'' % arg[0])

# Start log
#for i in range(1,len(sys.argv)):
#    checkArgs(sys.argv[i])

logging.info("### Starting new run ###")
logging.info("live is set to %s" % live)
logging.info("style is set to %s" % style)
logging.info("DYKMoverBot version %s" % version)

# Load the various pages
logging.info("Loading pages")
site = pywikibot.Site('en', 'wikipedia')
nomPage      = pywikibot.Page(site,'Template talk:Did you know')
#approvedPage = pywikibot.Page(site,'Template talk:Did you know/Approved')
approvedPage = pywikibot.Page(site,'User:Wugapodes/DYKTest/0')
approvedPage1= pywikibot.Page(site,'User:Wugapodes/DYKTest/1')

dateRegex = re.compile(r'on (.*?) (\d+)=')

DYKpage = nomPage.text.split('\n')
dates   = []
nonDate = []
entries = []
problem = []

# Get approved nominations on the DYK page
logging.info("Reading Nomination page")
for line in DYKpage:
    entries.append(line)
    if '==Articles' in line:
        matches=dateRegex.search(line)
        month = monthConvert(str(matches.group(1)))
        day = int(matches.group(2))
        dt = datetime.date(month=month,day=day,year=computeYear(month))
        section = line
        dates.append([dt,[section]])
    elif 'Did you know nominations/' in line and '<!--' not in line:
        if '}}{{' in line:
            logging.debug("Multiple nominations on one line.\n"\
                         +"Section: %s" % line)
            splitLine = line.split('}}{{')
            for title in splitLine:
                try:
                    checkPage(title)
                except Exception as e:
                    logging.error(
                        "Function checkPage did not complete successfully.\n"\
                        +str(e)+"\n"+title
                    )
                    logging.debug(line)
                    print(e)
                    problem.append([line,e])
                    continue
        else:
            line = line.split('}')[0]
            try:
                checkPage(line)
            except Exception as e:
                logging.error(
                    "Function checkPage did not complete successfully.\n"\
                    +str(e)+"\n"+line
                )
                print(e)
                problem.append([line,e])
                continue

# Get the text on the approved page and then update the proper sections
# with entries newly approved
logging.info("Reading Approved page")
approvedPageText = approvedPage.text.split('\n')
sectionName = ''
approvedPageDates = []
oldLine = ''
datesToRemove = []
if style != 1:
    logging.info("Starting style 0 parsing")
    for line in approvedPageText:
        if '==Articles' in line:
            mergeNominations(oldLine)
            matches=dateRegex.search(line)
            month = monthConvert(str(matches.group(1)))
            day = int(matches.group(2))
            dt = datetime.date(month=month,day=day,year=computeYear(month))
            sectionName = line
            for item in dates:
                if dt in item:
                    oldLine = dates.index(item)
                    datesToRemove.append(oldLine)
                    break
                else:
                    continue
            approvedPageDates.append([dt,[sectionName]])
        elif 'Did you know nominations/' in line and '<!--' not in line:
            if '}}{{' in line:
                splitLine = line.split('}}{{')
                for title in splitLine:
                    if computeNomStatus(title) >= 0:
                        try:
                            approvedPageDates[-1][1].append('{{'+title+'}}')
                        except Exception as e:
                            logging.warning("The approved page is empty?\n"+str(e))
            else:
                line = line.split('}')[0]
                if computeNomStatus(title) >= 0:
                    try:
                        approvedPageDates[-1][1].append('{{'+title+'}}')
                    except Exception as e:
                        logging.warning("The approved page is empty?\n"+str(e))
    dates2 = []
    for i in datesToRemove:
        dates2.append(dates[i])
    newSections = [x for x in dates if x not in dates2 and len(x[1]) > 1]
    approvedPageDates+=newSections
    toPrint=[]
    approvedPageDates.sort(key=lambda x: x[0])
    for entry in dates:
        if len(entry[1]) > 1:
            toPrint+=entry[1]

approvedPage1Text = approvedPage1.text.split('\n')
toRemoveFromNonDate = []
if style != 0:
    logging.info("Starting style 1 parsing.")
    logging.warning("Style 1 is no longer maintained. "\
        +"Bot may behave incorrectly or not at all.")
    for line in approvedPage1Text:
        if line in nonDate:
            toRemoveFromNonDate.append(line)
    nonDateTemp = [x for x in nonDate if x not in toRemoveFromNonDate]
    nonDate = nonDateTemp

# Create the page text to be output
logging.info("Creating output text")
approvedText2 = approvedPage1.text.split('\n')
passed = 0
approvedText = [
        [
        "{{/top}}\n",
        "=Nominations=\n",
        "==Approved nominations==\n",
        "<!-- This section will hold approved nominations, with the templates "\
        +"transcluded in the same manner as the regular nominations page. "\
        +"While the exact format of the section has not yet been decided-while"\
        +" it seems unlikely that it will be by date, it may be divided into "\
        +"other sections-it is likely that the oldest approvals will go at the"\
        +" top and the most recent ones at the bottom of each section. -->\n"
        ],
        [
        "{{/top}}\n",
        "=Nominations=\n",
        "==Approved nominations==\n",
        "<!-- This section will hold approved nominations, with the templates "\
        +"transcluded in the same manner as the regular nominations page. "\
        +"While the exact format of the section has not yet been decided-while"\
        +" it seems unlikely that it will be by date, it may be divided into "\
        +"other sections-it is likely that the oldest approvals will go at the"\
        +" top and the most recent ones at the bottom of each section. -->\n"
        ]
    ]

for line in approvedText2:
    if '==Approved nominations==' in line:
        passed=1
    elif '==Special occasion holding area==' in line:
        passed=2
        approvedText[1].append('\n'.join(nonDate))
        approvedText[1].append('\n'+line+'\n')
    elif passed > 0:
        approvedText[1].append(line+'\n')
approvedText[0].append('\n'.join(toPrint))
for line in approvedPage.text.split('\n'):
    if '==Special occasion holding area==' in line:
        passed = 1
        approvedText[0].append('\n'+line+'\n')
    elif passed == 1:
        approvedText[0].append(line+'\n')

# Determine if the bot should write to a live page or the test page. Defaults to
#     test page. Value of -1 tests backlog update (not standard because the file
#     size is very big).
if type(live) is str:
    pass
    page = pywikibot.Page(site,'Template talk:Did you know/Approved')
    page.text=''.join(approvedText[style])
    page.save('moving '+str(len(nonDate))+'tentatively approved nominations '\
              +'from [[WP:DYKN]], WugBot v'+version)
    page = pywikibot.Page(site,'Template talk:Did you know')
    # This is where the rest of the live stuff will go once I figure out
    #   how the other bots handle empty sections
else:
    page = pywikibot.Page(site,'User:Wugapodes/DYKTest/0')
    page.text=''.join(approvedText[0])
    page.save('test of DYKMover, WugBot v'+version+' section style')
    page = pywikibot.Page(site,'User:Wugapodes/DYKTest/1')
    page.text=''.join(approvedText[1])
    page.save('test of DYKMover, WugBot v'+version+' nosection style')
    page = pywikibot.Page(site,'User:Wugapodes/DYKTest/0')

print('Done')
<|endoftext|>


<commit_before>import threading
from queue import Queue

from netaddr import EUI, IPAddress
from scapy.all import *
from scapy.layers.inet import IP, TCP, ICMP
from scapy.layers.l2 import arping
from scapy.sendrecv import sr1

from analyzr.core import config
from analyzr.core.entities import NetworkNode
from analyzr.networkdiscovery.scanner import Scanner
from analyzr.topports import topports
from analyzr.utils.network import resolve_ip, TCPFlag, ScapyTCPFlag

logger = logging.getLogger(__name__)

class ArpPing(Scanner):
    def __init__(self):
        super(ArpPing, self).__init__()

    def scan(self):
        try:
            for interface, network in config.interfaces_networks.items():
                if not network.is_private():
                    self.logger.info(
                        "Skipping arp ping scan on network {0:s} because it's a public network".format(str(network)))
                    continue

                self.logger.info("Executing arp ping scan on network {0:s}...".format(str(network)))
                discovered_hosts = set()
                ans, unans = arping(str(network), iface=interface, timeout=1, verbose=False)
                for s, r in ans.res:
                    node = NetworkNode()
                    node.ip = IPAddress(r[ARP].psrc)
                    node.mac = EUI(r[Ether].src)
                    node.host = resolve_ip(r[ARP].psrc)
                    discovered_hosts.add(node)

                self.scan_results[network] = discovered_hosts
        except socket.error as e:
            if e.errno == socket.errno.EPERM:  # Operation not permitted
                self.logger.error("%s. Did you run as root?", e.strerror)
            else:
                raise

        self.logger.info("Arp ping scan done. Found %d unique hosts.", self.number_of_hosts_found)

class ICMPPing(Scanner):
    def __init__(self):
        super(ICMPPing, self).__init__()

    def scan(self):
        try:
            for interface, network in config.interfaces_networks.items():
                if network.is_private() and not config.scan_local_network_as_public:
                    self.logger.info(
                        "Skipping ICMP ping scan on network {0:s} because it's a private network.".format(str(network)))
                    continue

                self.logger.info("Executing ICMP ping scan on network {0:s}...".format(str(network)))
                discovered_hosts = set()

                ips_queue = Queue(network.size)
                for thread_id in range(config.num_ping_threads):
                    t = threading.Thread(
                        target=self._ping,
                        args=(ips_queue, interface, discovered_hosts,),
                        name='worker-{}'.format(thread_id),
                        daemon=True
                    )
                    t.start()

                # populate the queue
                [ips_queue.put(host_ip) for host_ip in network.iter_hosts()]

                self.logger.debug('*** main thread waiting')
                ips_queue.join()

                self.logger.debug(u"Got {0:d} answers.".format(len(discovered_hosts)))

                self.scan_results[network] = discovered_hosts
        except socket.error as e:
            if e.errno == socket.errno.EPERM:  # Operation not permitted
                self.logger.error("%s. Did you run as root?", e.strerror)
            else:
                raise

        self.logger.info("ICMP ping scan done. Found %d unique hosts.", self.number_of_hosts_found)

    def _ping(self, hosts: Queue, interface: str, results: set):
        self.logger.debug("{}: Starting ICMP ping thread.".format(threading.current_thread().name))
        while True:
            host = hosts.get()  # type: IPAddress
            host_str = str(host)

            res = sr1(IP(dst=host_str) / ICMP(), iface=interface, timeout=0.1, verbose=False)
            if res:
                node = NetworkNode()
                node.ip = host
                node.mac = EUI(res.src)
                node.host = resolve_ip(res[IP].src)
                results.add(node)

            hosts.task_done()

class TCPSYNPing(Scanner):
    def __init__(self):
        super(TCPSYNPing, self).__init__()
        if config.fastTCP:
            self.portstoscan = [80]
        else:
            self.portstoscan = topports

    def scan(self):
        """
        Attempts to identify if a host exists by sending a TCP SYN packet to a port. If we receive a SYN/ACK or a RST
        packet we know the host exists and we stop searching. This is not a port scanner. This does not find opened
        ports (it stops to search at the first sign that the host exists).
        """
        try:
            for interface, network in config.interfaces_networks.items():
                if network.is_private() and not config.scan_local_network_as_public:
                    self.logger.info(
                        "Skipping TCP ACK Ping on {0:s} because it's a private network.".format(str(network)))
                    continue

                if config.fastTCP:
                    self.logger.info("Executing TCP SYN ping scan (fast version) on {0:s}...".format(str(network)))
                else:
                    self.logger.info("Executing TCP SYN ping scan (slow version) on {0:s}...".format(str(network)))

                self.logger.info("Scanning ports %s.", str(self.portstoscan).strip("[]"))

                discovered_hosts = set()

                ips_queue = Queue(network.size)
                for thread_id in range(config.num_ping_threads):
                    t = threading.Thread(
                        target=self._port_ping,
                        args=(ips_queue, interface, discovered_hosts,),
                        name='worker-{}'.format(thread_id),
                        daemon=True
                    )
                    t.start()

                # populate the queue
                [ips_queue.put(host_ip) for host_ip in network.iter_hosts()]

                self.logger.debug('*** main thread waiting')
                ips_queue.join()

                self.logger.debug(u"Got {0:d} answers.".format(len(discovered_hosts)))

                self.scan_results[network] = discovered_hosts
        except socket.error as e:
            if e.errno == socket.errno.EPERM:  # Operation not permitted
                self.logger.error("%s. Did you run as root?", e.strerror)
            else:
                raise

        self.logger.info("TCP SYN ping scan done. Found %d unique hosts.", self.number_of_hosts_found)

    def _port_ping(self, hosts: Queue, interface: str, results: set):
        self.logger.debug("{}: Starting TCP SYN ping thread.".format(threading.current_thread().name))

        while True:
            host = hosts.get()  # type: IPAddress
            host_str = str(host)

            # Send SYN with random Src Port for each Dst port
            for dstPort in self.portstoscan:
                srcPort = random.randint(1025, 65534)
                resp = sr1(IP(dst=host_str) / TCP(sport=srcPort, dport=dstPort, flags=ScapyTCPFlag.SYN), timeout=1,
                           verbose=False,
                           iface=interface)
                if resp and resp.haslayer(TCP):
                    if resp[TCP].flags == (TCPFlag.SYN | TCPFlag.ACK) or resp[TCP].flags == (TCPFlag.RST | TCPFlag.ACK):
                        # Send Reset packet (RST)
                        send(IP(dst=host_str) / TCP(sport=srcPort, dport=dstPort, flags=ScapyTCPFlag.RST),
                             iface=interface, verbose=False)

                        # We know the port is closed or opened (we got a response), so we deduce that the host exists
                        node = NetworkNode()
                        node.ip = host
                        node.mac = EUI(resp.src)
                        node.host = resolve_ip(resp[IP].src)
                        results.add(node)

                        self.logger.debug(
                            "Found a live host by pinging port {port_nbr}: {live_host}.".format(port_nbr=dstPort,
                                                                                live_host=str(node)))

                        # We don't need to test the other ports. We know the host exists.
                        break

            hosts.task_done()
<commit_msg>Meilleur noms de variables à certains endroits.<commit_after>import threading
from queue import Queue

from netaddr import EUI, IPAddress
from scapy.all import *
from scapy.layers.inet import IP, TCP, ICMP
from scapy.layers.l2 import arping
from scapy.sendrecv import sr1

from analyzr.core import config
from analyzr.core.entities import NetworkNode
from analyzr.networkdiscovery.scanner import Scanner
from analyzr.topports import topports
from analyzr.utils.network import resolve_ip, TCPFlag, ScapyTCPFlag

logger = logging.getLogger(__name__)

class ArpPing(Scanner):
    def __init__(self):
        super(ArpPing, self).__init__()

    def scan(self):
        try:
            for interface, network in config.interfaces_networks.items():
                if not network.is_private():
                    self.logger.info(
                        "Skipping arp ping scan on network {0:s} because it's a public network".format(str(network)))
                    continue

                self.logger.info("Executing arp ping scan on network {0:s}...".format(str(network)))
                discovered_hosts = set()
                ans, unans = arping(str(network), iface=interface, timeout=1, verbose=False)
                for s, r in ans.res:
                    node = NetworkNode()
                    node.ip = IPAddress(r[ARP].psrc)
                    node.mac = EUI(r[Ether].src)
                    node.host = resolve_ip(r[ARP].psrc)
                    discovered_hosts.add(node)

                self.scan_results[network] = discovered_hosts
        except socket.error as e:
            if e.errno == socket.errno.EPERM:  # Operation not permitted
                self.logger.error("%s. Did you run as root?", e.strerror)
            else:
                raise

        self.logger.info("Arp ping scan done. Found %d unique hosts.", self.number_of_hosts_found)

class ICMPPing(Scanner):
    def __init__(self):
        super(ICMPPing, self).__init__()

    def scan(self):
        try:
            for interface, network in config.interfaces_networks.items():
                if network.is_private() and not config.scan_local_network_as_public:
                    self.logger.info(
                        "Skipping ICMP ping scan on network {0:s} because it's a private network.".format(str(network)))
                    continue

                self.logger.info("Executing ICMP ping scan on network {0:s}...".format(str(network)))
                discovered_hosts = set()

                ips_queue = Queue(network.size)
                for thread_id in range(config.num_ping_threads):
                    t = threading.Thread(
                        target=self._ping,
                        args=(ips_queue, interface, discovered_hosts,),
                        name='worker-{}'.format(thread_id),
                        daemon=True
                    )
                    t.start()

                # populate the queue
                [ips_queue.put(host_ip) for host_ip in network.iter_hosts()]

                self.logger.debug('*** main thread waiting')
                ips_queue.join()

                self.logger.debug(u"Got {0:d} answers.".format(len(discovered_hosts)))

                self.scan_results[network] = discovered_hosts
        except socket.error as e:
            if e.errno == socket.errno.EPERM:  # Operation not permitted
                self.logger.error("%s. Did you run as root?", e.strerror)
            else:
                raise

        self.logger.info("ICMP ping scan done. Found %d unique hosts.", self.number_of_hosts_found)

    def _ping(self, hosts: Queue, interface: str, results: set):
        self.logger.debug("{}: Starting ICMP ping thread.".format(threading.current_thread().name))
        while True:
            ip = hosts.get()  # type: IPAddress
            ip_str = str(ip)

            res = sr1(IP(dst=ip_str) / ICMP(), iface=interface, timeout=0.1, verbose=False)
            if res:
                node = NetworkNode()
                node.ip = ip
                node.mac = EUI(res.src)
                node.host = resolve_ip(res[IP].src)
                results.add(node)

            hosts.task_done()

class TCPSYNPing(Scanner):
    def __init__(self):
        super(TCPSYNPing, self).__init__()
        if config.fastTCP:
            self.portstoscan = [80]
        else:
            self.portstoscan = topports

    def scan(self):
        """
        Attempts to identify if a host exists by sending a TCP SYN packet to a port. If we receive a SYN/ACK or a RST
        packet we know the host exists and we stop searching. This is not a port scanner. This does not find opened
        ports (it stops to search at the first sign that the host exists).
        """
        try:
            for interface, network in config.interfaces_networks.items():
                if network.is_private() and not config.scan_local_network_as_public:
                    self.logger.info(
                        "Skipping TCP ACK Ping on {0:s} because it's a private network.".format(str(network)))
                    continue

                if config.fastTCP:
                    self.logger.info("Executing TCP SYN ping scan (fast version) on {0:s}...".format(str(network)))
                else:
                    self.logger.info("Executing TCP SYN ping scan (slow version) on {0:s}...".format(str(network)))

                self.logger.info("Scanning ports %s.", str(self.portstoscan).strip("[]"))

                discovered_hosts = set()

                ips_queue = Queue(network.size)
                for thread_nbr in range(config.num_ping_threads):
                    t = threading.Thread(
                        target=self._port_ping,
                        args=(ips_queue, interface, discovered_hosts,),
                        name='worker-{}'.format(thread_nbr),
                        daemon=True
                    )
                    t.start()

                # populate the queue
                [ips_queue.put(host_ip) for host_ip in network.iter_hosts()]

                self.logger.debug('*** main thread waiting')
                ips_queue.join()

                self.logger.debug(u"Got {0:d} answers.".format(len(discovered_hosts)))

                self.scan_results[network] = discovered_hosts
        except socket.error as e:
            if e.errno == socket.errno.EPERM:  # Operation not permitted
                self.logger.error("%s. Did you run as root?", e.strerror)
            else:
                raise

        self.logger.info("TCP SYN ping scan done. Found %d unique hosts.", self.number_of_hosts_found)

    def _port_ping(self, hosts: Queue, interface: str, results: set):
        self.logger.debug("{}: Starting TCP SYN ping thread.".format(threading.current_thread().name))

        while True:
            ip = hosts.get()  # type: IPAddress
            ip_str = str(ip)

            # Send SYN with random Src Port for each Dst port
            for dstPort in self.portstoscan:
                srcPort = random.randint(1025, 65534)
                resp = sr1(IP(dst=ip_str) / TCP(sport=srcPort, dport=dstPort, flags=ScapyTCPFlag.SYN), timeout=1,
                           verbose=False,
                           iface=interface)
                if resp and resp.haslayer(TCP):
                    if resp[TCP].flags == (TCPFlag.SYN | TCPFlag.ACK) or resp[TCP].flags == (TCPFlag.RST | TCPFlag.ACK):
                        # Send Reset packet (RST)
                        send(IP(dst=ip_str) / TCP(sport=srcPort, dport=dstPort, flags=ScapyTCPFlag.RST),
                             iface=interface, verbose=False)

                        # We know the port is closed or opened (we got a response), so we deduce that the host exists
                        node = NetworkNode()
                        node.ip = ip
                        node.mac = EUI(resp.src)
                        node.host = resolve_ip(resp[IP].src)
                        results.add(node)

                        self.logger.debug(
                            "Found a live host by pinging port {port_nbr}: {live_host}.".format(port_nbr=dstPort,
                                                                                live_host=str(node)))

                        # We don't need to test the other ports. We know the host exists.
                        break

            hosts.task_done()
<|endoftext|>


<commit_before>import logging
import collections

import networkx

from .function import Function

l = logging.getLogger(name="angr.knowledge.function_manager")

class FunctionDict(dict):
    """
    FunctionDict is a dict where the keys are function starting addresses and
    map to the associated :class:`Function`.
    """
    def __init__(self, backref, *args, **kwargs):
        self._backref = backref
        super(FunctionDict, self).__init__(*args, **kwargs)

    def __missing__(self, key):
        if isinstance(key, (int, long)):
            addr = key
        else:
            raise ValueError("FunctionDict.__missing__ only supports int as key type")

        t = Function(self._backref, addr)
        self[addr] = t
        return t

class FunctionManager(collections.Mapping):
    """
    This is a function boundaries management tool. It takes in intermediate
    results during CFG generation, and manages a function map of the binary.
    """
    def __init__(self, kb):
        self._kb = kb
        self._function_map = FunctionDict(self)
        self.callgraph = networkx.DiGraph()

        # Registers used for passing arguments around
        self._arg_registers = kb._project.arch.argument_registers

    def _genenare_callmap_sif(self, filepath):
        """
        Generate a sif file from the call map.

        :param filepath:    Path of the sif file
        :return:            None
        """
        with open(filepath, "wb") as f:
            for src, dst in self.callgraph.edges():
                f.write("%#x\tDirectEdge\t%#x\n" % (src, dst))

    def _add_node(self, function_addr, node, syscall=None):
        if type(node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            node = self._kb._project.factory.snippet(node)
        dst_func = self._function_map[function_addr]
        if syscall in (True, False):
            dst_func.is_syscall = syscall
        dst_func._register_nodes(node)

    def _add_call_to(self, function_addr, from_node, to_addr, retn_node, syscall=None):

        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(retn_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            retn_node = self._kb._project.factory.snippet(retn_node)
        dest_func = self._function_map[to_addr]
        if syscall in (True, False):
            dest_func.is_syscall = syscall

        self._function_map[function_addr]._call_to(from_node, dest_func, retn_node)
        self._function_map[function_addr]._add_call_site(from_node.addr, to_addr, retn_node.addr if retn_node else None)
        self.callgraph.add_edge(function_addr, to_addr)

    def _add_fakeret_to(self, function_addr, from_node, to_node, confirmed=None, syscall=None):
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        src_func = self._function_map[function_addr]

        if syscall in (True, False):
            src_func.is_syscall = syscall

        src_func._fakeret_to(from_node, to_node, confirmed=confirmed)

    def _remove_fakeret(self, function_addr, from_node, to_node):
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        self._function_map[function_addr]._remove_fakeret(from_node, to_node)

    def _add_return_from(self, function_addr, from_node, to_node=None): #pylint:disable=unused-argument
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        self._function_map[function_addr]._add_return_site(from_node)

    def _add_transition_to(self, function_addr, from_node, to_node):
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        self._function_map[function_addr]._transit_to(from_node, to_node)

    def _add_outside_transition_to(self, function_addr, from_node, to_node):
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        self._function_map[function_addr]._transit_to(from_node, to_node, outside=True)

    def _add_return_from_call(self, function_addr, src_function_addr, to_node):

        # Note that you will never return to a syscall

        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        self._function_map[function_addr]._return_from_call(
            self._function_map[src_function_addr], to_node)

    #
    # Dict methods
    #

    def __getitem__(self, k):
        if isinstance(k, (int, long)):
            f = self.function(addr=k)
        elif isinstance(k, str):
            f = self.function(name=k)
        else:
            raise ValueError("FunctionManager.__getitem__ deos not support keys of type %s" % type(k))

        if f is None:
            raise KeyError(k)

        return f

    def __setitem__(self, k, v):
        if isinstance(k, (int, long)):
            self._function_map[k] = v
        else:
            raise ValueError("FunctionManager.__setitem__ keys must be an int")

    def __delitem__(self, k):
        if isinstance(k, (int, long)):
            del self._function_map[k]
        else:
            raise ValueError("FunctionManager.__delitem__ only accepts int as key")

    def __len__(self):
        return len(self._function_map)

    def __iter__(self):
        for i in sorted(self._function_map.iterkeys()):
            yield i

    def function(self, addr=None, name=None, create=False, syscall=False):
        """
        Get a function object from the function manager.

        Pass either `addr` or `name` with the appropriate values.
        """
        if addr is not None:
            if addr in self._function_map:
                return self._function_map[addr]
            elif create:
                f = self._function_map[addr]
                if syscall:
                    f.is_syscall=True
                return f
        elif name is not None:
            for func in self._function_map.itervalues():
                if func.name == name:
                    return func

        return None

    def dbg_draw(self, prefix='dbg_function_'):
        for func_addr, func in self._function_map.iteritems():
            filename = "%s%#08x.png" % (prefix, func_addr)
            func.dbg_draw(filename)
<commit_msg>Implement copy() and clear() for FunctionManager<commit_after>import logging
import collections

import networkx

from .function import Function

l = logging.getLogger(name="angr.knowledge.function_manager")

class FunctionDict(dict):
    """
    FunctionDict is a dict where the keys are function starting addresses and
    map to the associated :class:`Function`.
    """
    def __init__(self, backref, *args, **kwargs):
        self._backref = backref
        super(FunctionDict, self).__init__(*args, **kwargs)

    def __missing__(self, key):
        if isinstance(key, (int, long)):
            addr = key
        else:
            raise ValueError("FunctionDict.__missing__ only supports int as key type")

        t = Function(self._backref, addr)
        self[addr] = t
        return t

class FunctionManager(collections.Mapping):
    """
    This is a function boundaries management tool. It takes in intermediate
    results during CFG generation, and manages a function map of the binary.
    """
    def __init__(self, kb):
        self._kb = kb
        self._function_map = FunctionDict(self)
        self.callgraph = networkx.DiGraph()

        # Registers used for passing arguments around
        self._arg_registers = kb._project.arch.argument_registers

    def copy(self):
        fm = FunctionManager(self._kb)
        fm._function_map = self._function_map.copy()
        fm.callgraph = networkx.DiGraph(self.callgraph)
        fm._arg_registers = self._arg_registers.copy()

        return fm

    def clear(self):
        self._function_map.clear()
        self.callgraph = networkx.DiGraph()

    def _genenare_callmap_sif(self, filepath):
        """
        Generate a sif file from the call map.

        :param filepath:    Path of the sif file
        :return:            None
        """
        with open(filepath, "wb") as f:
            for src, dst in self.callgraph.edges():
                f.write("%#x\tDirectEdge\t%#x\n" % (src, dst))

    def _add_node(self, function_addr, node, syscall=None):
        if type(node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            node = self._kb._project.factory.snippet(node)
        dst_func = self._function_map[function_addr]
        if syscall in (True, False):
            dst_func.is_syscall = syscall
        dst_func._register_nodes(node)

    def _add_call_to(self, function_addr, from_node, to_addr, retn_node, syscall=None):

        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(retn_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            retn_node = self._kb._project.factory.snippet(retn_node)
        dest_func = self._function_map[to_addr]
        if syscall in (True, False):
            dest_func.is_syscall = syscall

        self._function_map[function_addr]._call_to(from_node, dest_func, retn_node)
        self._function_map[function_addr]._add_call_site(from_node.addr, to_addr, retn_node.addr if retn_node else None)
        self.callgraph.add_edge(function_addr, to_addr)

    def _add_fakeret_to(self, function_addr, from_node, to_node, confirmed=None, syscall=None):
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        src_func = self._function_map[function_addr]

        if syscall in (True, False):
            src_func.is_syscall = syscall

        src_func._fakeret_to(from_node, to_node, confirmed=confirmed)

    def _remove_fakeret(self, function_addr, from_node, to_node):
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        self._function_map[function_addr]._remove_fakeret(from_node, to_node)

    def _add_return_from(self, function_addr, from_node, to_node=None): #pylint:disable=unused-argument
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        self._function_map[function_addr]._add_return_site(from_node)

    def _add_transition_to(self, function_addr, from_node, to_node):
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        self._function_map[function_addr]._transit_to(from_node, to_node)

    def _add_outside_transition_to(self, function_addr, from_node, to_node):
        if type(from_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            from_node = self._kb._project.factory.snippet(from_node)
        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        self._function_map[function_addr]._transit_to(from_node, to_node, outside=True)

    def _add_return_from_call(self, function_addr, src_function_addr, to_node):

        # Note that you will never return to a syscall

        if type(to_node) in (int, long):  # pylint: disable=unidiomatic-typecheck
            to_node = self._kb._project.factory.snippet(to_node)
        self._function_map[function_addr]._return_from_call(
            self._function_map[src_function_addr], to_node)

    #
    # Dict methods
    #

    def __getitem__(self, k):
        if isinstance(k, (int, long)):
            f = self.function(addr=k)
        elif isinstance(k, str):
            f = self.function(name=k)
        else:
            raise ValueError("FunctionManager.__getitem__ deos not support keys of type %s" % type(k))

        if f is None:
            raise KeyError(k)

        return f

    def __setitem__(self, k, v):
        if isinstance(k, (int, long)):
            self._function_map[k] = v
        else:
            raise ValueError("FunctionManager.__setitem__ keys must be an int")

    def __delitem__(self, k):
        if isinstance(k, (int, long)):
            del self._function_map[k]
        else:
            raise ValueError("FunctionManager.__delitem__ only accepts int as key")

    def __len__(self):
        return len(self._function_map)

    def __iter__(self):
        for i in sorted(self._function_map.iterkeys()):
            yield i

    def function(self, addr=None, name=None, create=False, syscall=False):
        """
        Get a function object from the function manager.

        Pass either `addr` or `name` with the appropriate values.
        """
        if addr is not None:
            if addr in self._function_map:
                return self._function_map[addr]
            elif create:
                f = self._function_map[addr]
                if syscall:
                    f.is_syscall=True
                return f
        elif name is not None:
            for func in self._function_map.itervalues():
                if func.name == name:
                    return func

        return None

    def dbg_draw(self, prefix='dbg_function_'):
        for func_addr, func in self._function_map.iteritems():
            filename = "%s%#08x.png" % (prefix, func_addr)
            func.dbg_draw(filename)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from __future__ import unicode_literals

"""
  Loading accessions statistics.
"""

import datetime, json, logging, pprint, re, sys, time, urllib
from datetime import date, timedelta
from optparse import make_option

import pymarc, requests
from django.contrib.auth.models import User
from django.core.cache import cache
from django.core.exceptions import ObjectDoesNotExist
from django.core.management.base import BaseCommand
from tech_services_reports import settings_app
from tech_services_reports.lib import marc_parser
from tech_services_reports.models import Accession, CatEdit, Harvest
from tech_services_reports.utility_code import CatStat
from tech_services_reports.utility_code import convert_date

log = logging.getLogger( "processing" )
log.debug( 'loaded ts_reports_loader.py' )

#For format mappings
url = settings_app.LOCATION_FORMAT_URL
requests.packages.urllib3.disable_warnings()
r = requests.get( url, verify=False )
location_format_map = json.loads( r.content.decode('utf-8') )
# log.debug( 'location_format_map, ```{}```'.format(pprint.pformat(location_format_map)) )
location_format_map = location_format_map['result']['items']

#For cat edits
#101111,sf,Catalog
# CAT_RE = re.compile('([0-9]{6})\,(\w+)\,(\w+)$')

TODAY = datetime.date.today()
#Number of days to go back for reports.
CUTOFF_DAY_DELTA = datetime.timedelta(days=2)

class Command(BaseCommand):
    help = "For loading Josiah exports for reports."

    def add_arguments(self, parser):
        parser.add_argument(
            '--summary', dest='summary', help='Loads from MARC exports.  Expecting items in MARC exports.'
            )
        parser.add_argument(
            '--cleaner',
            dest='cleaner',
            help='Cleans up database by finding common data entry problems.'
            )

    def __init__(self):
        self.last_harvest = self.last_harvest()

    def handle(self, **options):
        log.debug( 'options, ```{}```'.format(pprint.pformat(options)) )
        if options['summary']:
            self.summary(options['summary'])
        elif options['cleaner']:
            self.cleaner()
        else:
          print>>sys.stderr, 'no file specified'

    def first_item(self, items):
        """Get the first attached item and use that as the accessions
        count date.  E.g.  Items will be counted as an accessioned title
        on the date of the first attached item."""
        first_item = marc_parser.get_first_item( items )
        return first_item

    def counted_items(self):
        """ Grabs stored accession items from db.
            Called by Command.summary() """
        log.debug( 'starting counted_items()' )
        timestamp = datetime.datetime.now()
        numbers = cache.get('counted_items__numbers')
        if numbers is None:
            from tech_services_reports.models import Accession
            numbers = set([i.number for i in Accession.objects.all()])
            cache.set( 'counted_items__numbers', numbers, 60*60*24 )  # db grab cached for a day
        log.info( 'that took, `{}`'.format( str(datetime.datetime.now()-timestamp) ) )
        log.debug( 'first three numbers, ```{three}```; number-count, `{count}`'.format(three=list(numbers)[0:2], count=len(numbers)) )
        return numbers

    def last_harvest(self):
        from tech_services_reports.models import Harvest
        from tech_services_reports import settings_app
        from datetime import date, timedelta
        try:
            last = Harvest.objects.order_by('date')[0:1].get()
            return last
        except:
            hdate = date(settings_app.BEGIN_YEAR, settings_app.BEGIN_MONTH, 1)
            #Set last harvest to one day before starting harvest info.
            hdate = hdate - timedelta(days=1)
            obj, created = Harvest.objects.get_or_create(date=hdate)
            obj.note = 'Initial harvest'
            obj.save()
            return obj

    # def get_bib_created(self, this_record):
    #     from tech_services_reports.utility_code import convert_date
    #     d = this_record['907']['c']
    #     return convert_date(d)

    def get_bib_created( self, this_record ):
        dt = marc_parser.get_bib_created( this_record )
        return dt

    def summary(self, marc_file):
        """ Harvests data points from exported MARC fields.
            Date counts will include the date of a given harvest. """
        log.debug( 'starting summary()' )
        log.info( 'reading MARC file, ```{}```'.format(marc_file) )

        #Dicts to store counts
        cataloging_edit_count = {}
        # cataloging_count = {}
        title_count = {}
        volume_count = {}

        #Find items already counted.
        #Add logic to skip counted items.
        log.info( 'retrieving existing items stored in Accessions database' )
        existing_items = self.counted_items()

        #======================================================================
        # parse marc file
        #======================================================================

        #Loop through marc records.
        # ( cataloging_edit_count, title_count, volume_count ) = marc_parser.parse_marc_file( marc_file, existing_items )
        from tech_services_reports.lib.marc_parser import FileParser
        file_parser = FileParser()
        ( cataloging_edit_count, title_count, volume_count ) = file_parser.process_marc_file( marc_filepath, existing_items, location_format_map )

        log.debug( 'cataloging_edit_count, ```{}```'.format( pprint.pformat(cataloging_edit_count) ) )
        log.debug( 'title_count, ```{}```'.format( pprint.pformat(title_count) ) )
        log.debug( 'volume_count dct, ```{}```'.format( pprint.pformat(volume_count) ) )

        #Finish looping through MARC records

        #======================================================================
        # put individual counts together into  db fixture
        #======================================================================

        db_fixture = []
        # print>>sys.stderr, "Creating accessions stats fixture."
        log.info( 'creating accessions stats fixture' )
        for meta, count in volume_count.items():
            #Prep fixture.
            fixd = {}
            fixd['model'] = "tech_services_reports.accession"
            fixd['pk'] = meta.number
            f = {}
            #f['number'] = meta.number
            f['created'] = meta.created.isoformat()
            f['acquisition_method'] = meta.acquisition_method
            f['format'] = meta.format
            f['location'] = meta.location
            f['volumes'] = count
            try:
                title_val = title_count[meta]
            except KeyError:
                title_val = 0
            f['titles'] = title_val

            f['serial_added_volume'] = meta.serial_added_volume

            fixd['fields'] = f

            db_fixture.append(fixd)

        # print>>sys.stderr, "Creating cataloging stats fixtures."
        log.info( 'creating cataloging stats fixture' )
        for meta, count in cataloging_edit_count.items():
            fixd = {}
            fixd['model'] = "tech_services_reports.catedit"
            #Create a unique stamp to serve as the primary key.
            #Important not to change this as the database moves
            #forward.
            cataloger, edate, ctype, bib, mat_type, source = meta
            meta_str = ''.join([bib,
                                cataloger,
                                ctype,
                                mat_type,
                                source,
                                edate.isoformat(),])
            if len(meta_str) > 50:
                meta_str = meta_str.replace('Backstage Library Works', 'BL')
            if len(meta_str) > 50:
                # print>>sys.stderr, "Can't record edit for {0}.  Edit string too long.".format(bib)
                log.warning( "can't record edit for bib `{}`; edit string too long".format(bib) )
                continue
            fixd['pk'] = meta_str
            f = {}
            f['editor'] = cataloger
            f['edit_date'] = edate.isoformat()
            f['type'] = ctype
            f['bib'] = bib
            f['mat_type'] = mat_type
            f['source'] = source
            fixd['fields'] = f
            db_fixture.append(fixd)

        #Done.  Print the fixture.
        # print simplejson.dumps(db_fixture)
        print( json.dumps(db_fixture) )

        ## end def summary()

    def count_volumes( self, marc_items, cat_date, material_type, counted_items ):
        """Create summary accession info for items created within given range."""
        volumes_and_titles_dct = marc_parser.count_volumes( marc_items, cat_date, material_type, counted_items )
        return volumes_and_titles_dct

    def count_cataloging_edits(self,
                               bib_number,
                               mat_type,
                               marc_995,
                               cat_edit_count,
                               source):
        cat_edit_count = marc_parser.count_cataloging_edits(
            bib_number, mat_type, marc_995, at_edit_count, source )
        return cat_edit_count

    def count_batch_edits(self,
                           bib_number,
                           create_date,
                           mat_type,
                           marc_910,
                           cat_edit_count,
                           source):
        """Count Shelf Ready bibs.  Email from Sam.
        YBP (b56908994) and Aux Amateurs (b57122490) use constant data in 910 $a.
        For Casalini (b57128510), the constant data we're interested in is
        placed in 910 $g.
        Also note that records created even a few weeks ago have overlays from our authority control vendor, Backstage Library Works; the YBP record also has an overlay from OCLC BibNote.  In these cases, the original cat date -- now lost -- equals the create date."""
        cat_edit_count = marc_parser.count_batch_edits(
            bib_number, create_date, mat_type, marc_910, cat_edit_count, source )
        return cat_edit_count

    def cleaner(self):
        """
        Utility to clean up common data entry problems.
        """
        #from tech_services_reports.utility_code import convert_date, CatStat
        from tech_services_reports.models import Accession, CatEdit, Harvest
        from datetime import date

        # location_format_map = simplejson.load(urllib.urlopen(settings_app.LOCATION_FORMAT_URL))
        location_format_map = json.loads(urllib.urlopen(settings_app.LOCATION_FORMAT_URL))
        location_format_map = location_format_map['result']['items']

        edits = CatEdit.objects.filter(type="Batch Load")
        print>>sys.stderr, "Will merge %s Batch Loads to Batch." % len(edits)
        for e in edits:
            print>>sys.stderr, e
            e.type = 'Batch'
            e.save()

        #Try to re-map unknown formats in accessions.
        unknowns = Accession.objects.filter(format="unknown")
        print>>sys.stderr, "Removing known unknowns -- approximately %s." % len(unknowns)
        for acc in unknowns:
            loc_code = str(acc.location)
            try:
                format = location_format_map[loc_code]['format']
                #print loc_code, format
                acc.format = format
                acc.save()
            except KeyError as e:
                print>>sys.stderr, "Bad code: ", loc_code

#        edits = CatEdit.objects.filter(type="Batch Load")
#        print>>sys.stderr, "Will merge %s Batch Loads to Batch." % len(edits)
#        for e in edits:
#            print>>sys.stderr, e
#            e.type = 'Batch'
#            e.save()

<commit_msg>upates var name.<commit_after># -*- coding: utf-8 -*-

from __future__ import unicode_literals

"""
  Loading accessions statistics.
"""

import datetime, json, logging, pprint, re, sys, time, urllib
from datetime import date, timedelta
from optparse import make_option

import pymarc, requests
from django.contrib.auth.models import User
from django.core.cache import cache
from django.core.exceptions import ObjectDoesNotExist
from django.core.management.base import BaseCommand
from tech_services_reports import settings_app
from tech_services_reports.lib import marc_parser
from tech_services_reports.models import Accession, CatEdit, Harvest
from tech_services_reports.utility_code import CatStat
from tech_services_reports.utility_code import convert_date

log = logging.getLogger( "processing" )
log.debug( 'loaded ts_reports_loader.py' )

#For format mappings
url = settings_app.LOCATION_FORMAT_URL
requests.packages.urllib3.disable_warnings()
r = requests.get( url, verify=False )
location_format_map = json.loads( r.content.decode('utf-8') )
# log.debug( 'location_format_map, ```{}```'.format(pprint.pformat(location_format_map)) )
location_format_map = location_format_map['result']['items']

#For cat edits
#101111,sf,Catalog
# CAT_RE = re.compile('([0-9]{6})\,(\w+)\,(\w+)$')

TODAY = datetime.date.today()
#Number of days to go back for reports.
CUTOFF_DAY_DELTA = datetime.timedelta(days=2)

class Command(BaseCommand):
    help = "For loading Josiah exports for reports."

    def add_arguments(self, parser):
        parser.add_argument(
            '--summary', dest='summary', help='Loads from MARC exports.  Expecting items in MARC exports.'
            )
        parser.add_argument(
            '--cleaner',
            dest='cleaner',
            help='Cleans up database by finding common data entry problems.'
            )

    def __init__(self):
        self.last_harvest = self.last_harvest()

    def handle(self, **options):
        log.debug( 'options, ```{}```'.format(pprint.pformat(options)) )
        if options['summary']:
            self.summary(options['summary'])
        elif options['cleaner']:
            self.cleaner()
        else:
          print>>sys.stderr, 'no file specified'

    def first_item(self, items):
        """Get the first attached item and use that as the accessions
        count date.  E.g.  Items will be counted as an accessioned title
        on the date of the first attached item."""
        first_item = marc_parser.get_first_item( items )
        return first_item

    def counted_items(self):
        """ Grabs stored accession items from db.
            Called by Command.summary() """
        log.debug( 'starting counted_items()' )
        timestamp = datetime.datetime.now()
        numbers = cache.get('counted_items__numbers')
        if numbers is None:
            from tech_services_reports.models import Accession
            numbers = set([i.number for i in Accession.objects.all()])
            cache.set( 'counted_items__numbers', numbers, 60*60*24 )  # db grab cached for a day
        log.info( 'that took, `{}`'.format( str(datetime.datetime.now()-timestamp) ) )
        log.debug( 'first three numbers, ```{three}```; number-count, `{count}`'.format(three=list(numbers)[0:2], count=len(numbers)) )
        return numbers

    def last_harvest(self):
        from tech_services_reports.models import Harvest
        from tech_services_reports import settings_app
        from datetime import date, timedelta
        try:
            last = Harvest.objects.order_by('date')[0:1].get()
            return last
        except:
            hdate = date(settings_app.BEGIN_YEAR, settings_app.BEGIN_MONTH, 1)
            #Set last harvest to one day before starting harvest info.
            hdate = hdate - timedelta(days=1)
            obj, created = Harvest.objects.get_or_create(date=hdate)
            obj.note = 'Initial harvest'
            obj.save()
            return obj

    # def get_bib_created(self, this_record):
    #     from tech_services_reports.utility_code import convert_date
    #     d = this_record['907']['c']
    #     return convert_date(d)

    def get_bib_created( self, this_record ):
        dt = marc_parser.get_bib_created( this_record )
        return dt

    def summary(self, marc_filepath):
        """ Harvests data points from exported MARC fields.
            Date counts will include the date of a given harvest. """
        log.debug( 'starting summary()' )
        log.info( 'reading MARC file, ```{}```'.format(marc_filepath) )

        #Dicts to store counts
        cataloging_edit_count = {}
        # cataloging_count = {}
        title_count = {}
        volume_count = {}

        #Find items already counted.
        #Add logic to skip counted items.
        log.info( 'retrieving existing items stored in Accessions database' )
        existing_items = self.counted_items()

        #======================================================================
        # parse marc file
        #======================================================================

        #Loop through marc records.
        # ( cataloging_edit_count, title_count, volume_count ) = marc_parser.parse_marc_file( marc_file, existing_items )
        from tech_services_reports.lib.marc_parser import FileParser
        file_parser = FileParser()
        ( cataloging_edit_count, title_count, volume_count ) = file_parser.process_marc_file( marc_filepath, existing_items, location_format_map )

        log.debug( 'cataloging_edit_count, ```{}```'.format( pprint.pformat(cataloging_edit_count) ) )
        log.debug( 'title_count, ```{}```'.format( pprint.pformat(title_count) ) )
        log.debug( 'volume_count dct, ```{}```'.format( pprint.pformat(volume_count) ) )

        #Finish looping through MARC records

        #======================================================================
        # put individual counts together into  db fixture
        #======================================================================

        db_fixture = []
        # print>>sys.stderr, "Creating accessions stats fixture."
        log.info( 'creating accessions stats fixture' )
        for meta, count in volume_count.items():
            #Prep fixture.
            fixd = {}
            fixd['model'] = "tech_services_reports.accession"
            fixd['pk'] = meta.number
            f = {}
            #f['number'] = meta.number
            f['created'] = meta.created.isoformat()
            f['acquisition_method'] = meta.acquisition_method
            f['format'] = meta.format
            f['location'] = meta.location
            f['volumes'] = count
            try:
                title_val = title_count[meta]
            except KeyError:
                title_val = 0
            f['titles'] = title_val

            f['serial_added_volume'] = meta.serial_added_volume

            fixd['fields'] = f

            db_fixture.append(fixd)

        # print>>sys.stderr, "Creating cataloging stats fixtures."
        log.info( 'creating cataloging stats fixture' )
        for meta, count in cataloging_edit_count.items():
            fixd = {}
            fixd['model'] = "tech_services_reports.catedit"
            #Create a unique stamp to serve as the primary key.
            #Important not to change this as the database moves
            #forward.
            cataloger, edate, ctype, bib, mat_type, source = meta
            meta_str = ''.join([bib,
                                cataloger,
                                ctype,
                                mat_type,
                                source,
                                edate.isoformat(),])
            if len(meta_str) > 50:
                meta_str = meta_str.replace('Backstage Library Works', 'BL')
            if len(meta_str) > 50:
                # print>>sys.stderr, "Can't record edit for {0}.  Edit string too long.".format(bib)
                log.warning( "can't record edit for bib `{}`; edit string too long".format(bib) )
                continue
            fixd['pk'] = meta_str
            f = {}
            f['editor'] = cataloger
            f['edit_date'] = edate.isoformat()
            f['type'] = ctype
            f['bib'] = bib
            f['mat_type'] = mat_type
            f['source'] = source
            fixd['fields'] = f
            db_fixture.append(fixd)

        #Done.  Print the fixture.
        # print simplejson.dumps(db_fixture)
        print( json.dumps(db_fixture) )

        ## end def summary()

    def count_volumes( self, marc_items, cat_date, material_type, counted_items ):
        """Create summary accession info for items created within given range."""
        volumes_and_titles_dct = marc_parser.count_volumes( marc_items, cat_date, material_type, counted_items )
        return volumes_and_titles_dct

    def count_cataloging_edits(self,
                               bib_number,
                               mat_type,
                               marc_995,
                               cat_edit_count,
                               source):
        cat_edit_count = marc_parser.count_cataloging_edits(
            bib_number, mat_type, marc_995, at_edit_count, source )
        return cat_edit_count

    def count_batch_edits(self,
                           bib_number,
                           create_date,
                           mat_type,
                           marc_910,
                           cat_edit_count,
                           source):
        """Count Shelf Ready bibs.  Email from Sam.
        YBP (b56908994) and Aux Amateurs (b57122490) use constant data in 910 $a.
        For Casalini (b57128510), the constant data we're interested in is
        placed in 910 $g.
        Also note that records created even a few weeks ago have overlays from our authority control vendor, Backstage Library Works; the YBP record also has an overlay from OCLC BibNote.  In these cases, the original cat date -- now lost -- equals the create date."""
        cat_edit_count = marc_parser.count_batch_edits(
            bib_number, create_date, mat_type, marc_910, cat_edit_count, source )
        return cat_edit_count

    def cleaner(self):
        """
        Utility to clean up common data entry problems.
        """
        #from tech_services_reports.utility_code import convert_date, CatStat
        from tech_services_reports.models import Accession, CatEdit, Harvest
        from datetime import date

        # location_format_map = simplejson.load(urllib.urlopen(settings_app.LOCATION_FORMAT_URL))
        location_format_map = json.loads(urllib.urlopen(settings_app.LOCATION_FORMAT_URL))
        location_format_map = location_format_map['result']['items']

        edits = CatEdit.objects.filter(type="Batch Load")
        print>>sys.stderr, "Will merge %s Batch Loads to Batch." % len(edits)
        for e in edits:
            print>>sys.stderr, e
            e.type = 'Batch'
            e.save()

        #Try to re-map unknown formats in accessions.
        unknowns = Accession.objects.filter(format="unknown")
        print>>sys.stderr, "Removing known unknowns -- approximately %s." % len(unknowns)
        for acc in unknowns:
            loc_code = str(acc.location)
            try:
                format = location_format_map[loc_code]['format']
                #print loc_code, format
                acc.format = format
                acc.save()
            except KeyError as e:
                print>>sys.stderr, "Bad code: ", loc_code

#        edits = CatEdit.objects.filter(type="Batch Load")
#        print>>sys.stderr, "Will merge %s Batch Loads to Batch." % len(edits)
#        for e in edits:
#            print>>sys.stderr, e
#            e.type = 'Batch'
#            e.save()

<|endoftext|>


<commit_before>from __future__ import print_function, division
import numpy as np
from ciabatta.meta import make_repr_str
import ahoy
from ahoy import ships, obstructors, fields, turners

class Model(object):

    def __init__(self, seed, dt,
                 aligned_flag=False, origin_flags=None,
                 **ship_kwargs):
        # Record initial conditions
        self.seed = seed
        self.dt = dt

        # Initialise system
        self.i = 0
        self.rng = np.random.RandomState(seed)
        self.ships = ships.ships_factory(self.rng,
                                         aligned_flag=aligned_flag,
                                         origin_flags=origin_flags,
                                         **ship_kwargs)

    @property
    def aligned_flags(self):
        return np.all(self.ships.agents.directions.u_0 == 0.0, axis=0)

    @property
    def origin_flags(self):
        return np.all(self.ships.agents.positions.r_0 == 0.0, axis=0)

    def iterate(self):
        self.ships.iterate(self.dt, self.rng)
        self.i += 1

    def _get_output_dirname_agent_part(self):
        ags = self.ships.agents

        origin_str = '({})'.format(','.join(['{:d}'.format(a)
                                             for a in self.aligned_flags]))
        s = 'n={},align={}'.format(ags.positions.n, origin_str)

        # Space and swimming
        if isinstance(ags.positions, ahoy.positions.Positions):
            origin_str = '({})'.format(','.join(['{:d}'.format(o)
                                                 for o in self.origin_flags]))
            s += ',origin={},v={:g}'.format(origin_str, ags.swimmers.v_0)
            if isinstance(ags.positions, ahoy.positions.PeriodicPositions):
                def format_inf(x):
                    return '{:g}'.format(x) if np.isfinite(x) else 'i'
                L_str = '({})'.format(','.join([format_inf(e)
                                                for e in ags.positions.L]))
                s += ',L={}'.format(L_str)

        # Rudders
        for rs in ags.rudder_sets.sets:
            nm = rs.noise_measurer
            if isinstance(rs, ahoy.rudders.TumbleRudders):
                noise_str = 'p'
            elif isinstance(rs, ahoy.rudders.RotationRudders):
                noise_str = 'Dr'
            s += ',{}={:g}'.format(noise_str, nm.noise_0)
            if rs.is_chemotactic:
                type_s = 'T' if nm.is_temporal else 'S'
                side = 2 - rs.is_onesided
                s += ',chi={:.2g},side={:d},type={}'.format(nm.chi, side,
                                                           type_s)
                if nm.is_temporal:
                    measurer = nm.dc_dx_measurer
                    s += ',dtMem={:g},tMem={:g}'.format(measurer.dt_mem,
                                                        measurer.t_mem)
        return s

    def _get_output_dirname_obstruction_part(self):
        obs = self.ships.obstructor
        s = 'obs='

        if obs.__class__ is obstructors.NoneObstructor:
            s += 'NoObs'
        elif obs.__class__ is obstructors.PorousObstructor:
            if obs.turner.__class__ is turners.Turner:
                s_turner = 'stall'
            elif obs.turner.__class__ is turners.BounceBackTurner:
                s_turner = 'bback'
            elif obs.turner.__class__ is turners.ReflectTurner:
                s_turner = 'reflect'
            elif obs.turner.__class__ is turners.AlignTurner:
                s_turner = 'align'
            pf = obs.fraction_occupied
            s += 'Pore(R={:g},pf={:.2g},turn={})'.format(obs.R, pf, s_turner)
        return s

    def _get_output_dirname_field_part(self):
        c_field = self.ships.c_field
        s = 'c='
        if c_field.__class__ is fields.NoneFoodField:
            s += 'NoC'
        elif c_field.__class__ is fields.FoodField:
            s += 'C(c0={:g},cD={:g},cDelta={:g})'.format(c_field.c_0,
                                                         c_field.D,
                                                         c_field.delta)
        return s

    def get_output_dirname(self):
        s = 'ahoy_{}D,dt={:g},seed={}'.format(self.ships.dim, self.dt,
                                              self.seed)
        s += ',{}'.format(self._get_output_dirname_agent_part())
        s += ',{}'.format(self._get_output_dirname_obstruction_part())
        s += ',{}'.format(self._get_output_dirname_field_part())
        return s

    def __repr__(self):
        fs = [('seed', self.seed), ('dt', self.dt),
              ('origin_flags', self.origin_flags),
              ('aligned_flags', self.aligned_flags),
              ('i', self.i), ('rng', self.rng),
              ('ships', self.ships)
              ]
        return make_repr_str(self, fs)
<commit_msg>Fix bad variable label when making model output directory name<commit_after>from __future__ import print_function, division
import numpy as np
from ciabatta.meta import make_repr_str
import ahoy
from ahoy import ships, obstructors, fields, turners

class Model(object):

    def __init__(self, seed, dt,
                 aligned_flag=False, origin_flags=None,
                 **ship_kwargs):
        # Record initial conditions
        self.seed = seed
        self.dt = dt

        # Initialise system
        self.i = 0
        self.rng = np.random.RandomState(seed)
        self.ships = ships.ships_factory(self.rng,
                                         aligned_flag=aligned_flag,
                                         origin_flags=origin_flags,
                                         **ship_kwargs)

    @property
    def aligned_flags(self):
        return np.all(self.ships.agents.directions.u_0 == 0.0, axis=0)

    @property
    def origin_flags(self):
        return np.all(self.ships.agents.positions.r_0 == 0.0, axis=0)

    def iterate(self):
        self.ships.iterate(self.dt, self.rng)
        self.i += 1

    def _get_output_dirname_agent_part(self):
        ags = self.ships.agents

        aligned_str = '({})'.format(','.join(['{:d}'.format(a)
                                              for a in self.aligned_flags]))
        s = 'n={},align={}'.format(ags.positions.n, aligned_str)

        # Space and swimming
        if isinstance(ags.positions, ahoy.positions.Positions):
            origin_str = '({})'.format(','.join(['{:d}'.format(o)
                                                 for o in self.origin_flags]))
            s += ',origin={},v={:g}'.format(origin_str, ags.swimmers.v_0)
            if isinstance(ags.positions, ahoy.positions.PeriodicPositions):
                def format_inf(x):
                    return '{:g}'.format(x) if np.isfinite(x) else 'i'
                L_str = '({})'.format(','.join([format_inf(e)
                                                for e in ags.positions.L]))
                s += ',L={}'.format(L_str)

        # Rudders
        for rs in ags.rudder_sets.sets:
            nm = rs.noise_measurer
            if isinstance(rs, ahoy.rudders.TumbleRudders):
                noise_str = 'p'
            elif isinstance(rs, ahoy.rudders.RotationRudders):
                noise_str = 'Dr'
            s += ',{}={:g}'.format(noise_str, nm.noise_0)
            if rs.is_chemotactic:
                type_s = 'T' if nm.is_temporal else 'S'
                side = 2 - rs.is_onesided
                s += ',chi={:.2g},side={:d},type={}'.format(nm.chi, side,
                                                           type_s)
                if nm.is_temporal:
                    measurer = nm.dc_dx_measurer
                    s += ',dtMem={:g},tMem={:g}'.format(measurer.dt_mem,
                                                        measurer.t_mem)
        return s

    def _get_output_dirname_obstruction_part(self):
        obs = self.ships.obstructor
        s = 'obs='

        if obs.__class__ is obstructors.NoneObstructor:
            s += 'NoObs'
        elif obs.__class__ is obstructors.PorousObstructor:
            if obs.turner.__class__ is turners.Turner:
                s_turner = 'stall'
            elif obs.turner.__class__ is turners.BounceBackTurner:
                s_turner = 'bback'
            elif obs.turner.__class__ is turners.ReflectTurner:
                s_turner = 'reflect'
            elif obs.turner.__class__ is turners.AlignTurner:
                s_turner = 'align'
            pf = obs.fraction_occupied
            s += 'Pore(R={:g},pf={:.2g},turn={})'.format(obs.R, pf, s_turner)
        return s

    def _get_output_dirname_field_part(self):
        c_field = self.ships.c_field
        s = 'c='
        if c_field.__class__ is fields.NoneFoodField:
            s += 'NoC'
        elif c_field.__class__ is fields.FoodField:
            s += 'C(c0={:g},cD={:g},cDelta={:g})'.format(c_field.c_0,
                                                         c_field.D,
                                                         c_field.delta)
        return s

    def get_output_dirname(self):
        s = 'ahoy_{}D,dt={:g},seed={}'.format(self.ships.dim, self.dt,
                                              self.seed)
        s += ',{}'.format(self._get_output_dirname_agent_part())
        s += ',{}'.format(self._get_output_dirname_obstruction_part())
        s += ',{}'.format(self._get_output_dirname_field_part())
        return s

    def __repr__(self):
        fs = [('seed', self.seed), ('dt', self.dt),
              ('origin_flags', self.origin_flags),
              ('aligned_flags', self.aligned_flags),
              ('i', self.i), ('rng', self.rng),
              ('ships', self.ships)
              ]
        return make_repr_str(self, fs)
<|endoftext|>