<commit_before>import unittest
from array import array

import ROOT
from libcppyy import SetOwnership, addressof
import numpy as np

class TTreeBranch(unittest.TestCase):
    """
    Test for the pythonization of TTree::Branch, which allows to pass proxy
    references as arguments from the Python side. Example:
    `v = ROOT.std.vector('int')()`
    `t.Branch('my_vector_branch', v)`
    """

    filename  = 'treebranch.root'
    treename  = 'mytree'
    arraysize = 10
    ival      = 7
    fval      = 7.

    # Setup
    @classmethod
    def setUpClass(cls):
        ROOT.gInterpreter.Declare("""
        struct MyStruct {
            int myint1;
            int myint2;
        };
        """)

    # Helpers
    def create_file_and_tree(self):
        f = ROOT.TFile(self.filename, 'RECREATE')
        t = ROOT.TTree(self.treename, self.treename)
        # Prevent double deletion of the tree (Python and C++ TFile)
        SetOwnership(t, False)

        return f,t

    def get_tree(self):
        f = ROOT.TFile(self.filename)
        t = f.Get(self.treename)
        SetOwnership(t, False)

        return f,t

    @staticmethod
    def fill_and_close(f, t):
        t.Fill()
        f.Write()
        f.Close()

    # Tests
    # Basic type and array do not actually need the pythonization,
    # but testing anyway for the sake of completeness
    def test01_write_basic_type_branch(self):
        f,t = self.create_file_and_tree()

        n = array('f', [ self.fval ])
        t.Branch('floatb', n, 'floatb/F')

        self.fill_and_close(f, t)

    def test02_read_basic_type_branch(self):
        f,t = self.get_tree()

        for entry in t:
            self.assertEqual(entry.floatb, self.fval)

    def test03_write_array_branch(self):
        f,t = self.create_file_and_tree()

        a = array('d', self.arraysize*[ self.fval ])
        t.Branch('arrayb', a, 'arrayb[' + str(self.arraysize) + ']/D')

        self.fill_and_close(f, t)

    def test04_read_array_branch(self):
        f,t = self.get_tree()

        for entry in t:
            for elem in entry.arrayb:
                self.assertEqual(elem, self.fval)

    def test05_write_numpy_array_branch(self):
        f,t = self.create_file_and_tree()

        a = np.array(self.arraysize*[ self.fval ]) # dtype='float64'
        t.Branch('arrayb', a, 'arrayb[' + str(self.arraysize) + ']/D')

        self.fill_and_close(f, t)

    def test06_read_numpy_array_branch(self):
        f,t = self.get_tree()

        for entry in t:
            for elem in entry.arrayb:
                self.assertEqual(elem, self.fval)

    # Struct and vector do benefit from the pythonization
    def test07_write_struct_single_branch(self):
        f,t = self.create_file_and_tree()

        ms = ROOT.MyStruct()
        ms.myint1, ms.myint2 = self.ival, 2*self.ival

        # Test overloads
        t.Branch('structb0', ms)
        t.Branch('structb1', ms, 32000)
        t.Branch('structb2', ms, 32000, 99)
        t.Branch('structb3', 'MyStruct', ms)
        t.Branch('structb4', 'MyStruct', ms, 32000)
        t.Branch('structb5', 'MyStruct', ms, 32000, 99)

        self.fill_and_close(f, t)

    def test08_read_struct_single_branch(self):
        f,t = self.get_tree()

        for entry in t:
            for ms in [ getattr(entry, 'structb' + str(i)) for i in range(6) ]:
                self.assertEqual(ms.myint1, self.ival)
                self.assertEqual(ms.myint2, 2*self.ival)

    def test09_write_struct_separate_branches(self):
        f,t = self.create_file_and_tree()

        ms = ROOT.MyStruct()
        ms.myint1, ms.myint2 = self.ival, 2*self.ival

        # Use `addressof` to get the address of the struct members
        t.Branch('myint1b', addressof(ms, 'myint1'), 'myint1balias/I')
        t.Branch('myint2b', addressof(ms, 'myint2'), 'myint2balias/I')

        self.fill_and_close(f, t)

    def test10_read_struct_separate_branches(self):
        f,t = self.get_tree()

        for entry in t:
            self.assertEqual(entry.myint1b, self.ival)
            self.assertEqual(entry.myint2b, 2*self.ival)
            # Test aliases
            self.assertEqual(entry.myint1balias, self.ival)
            self.assertEqual(entry.myint2balias, 2*self.ival)

    def test11_write_vector_branch(self):
        f,t = self.create_file_and_tree()

        v = ROOT.std.vector('double')(self.arraysize*[ self.fval ])

        # Test overloads
        t.Branch('vectorb0', v)
        t.Branch('vectorb1', v, 32000)
        t.Branch('vectorb2', v, 32000, 99)
        t.Branch('vectorb3', 'std::vector<double>', v)
        t.Branch('vectorb4', 'std::vector<double>', v, 32000)
        t.Branch('vectorb5', 'std::vector<double>', v, 32000, 99)

        self.fill_and_close(f, t)

    def test12_read_vector_branch(self):
        f,t = self.get_tree()

        for entry in t:
            for v in [ getattr(entry, 'vectorb' + str(i)) for i in range(6) ]:
                for elem in v:
                    self.assertEqual(elem, self.fval)

    def test13_write_fallback_case(self):
        f,t = self.create_file_and_tree()

        # Test an overload that uses the original Branch proxy
        l = ROOT.TList()
        s = ROOT.TString('one:two')
        a = s.Tokenize(ROOT.TString(':'))
        a.SetName('myobjarray')
        l.Add(a)
        t.Branch(l)

        self.fill_and_close(f, t)

    def test14_read_fallback_case(self):
        f,t = self.get_tree()

        for entry in t:
            self.assertEqual(entry.myobjarray_one.GetName(), 'one')
            self.assertEqual(entry.myobjarray_two.GetName(), 'two')

if __name__ == '__main__':
    unittest.main()
<commit_msg>[Exp PyROOT] SetOwnership and addressof can be imported from ROOT<commit_after>import unittest
from array import array

import ROOT
from ROOT import SetOwnership, addressof
import numpy as np

class TTreeBranch(unittest.TestCase):
    """
    Test for the pythonization of TTree::Branch, which allows to pass proxy
    references as arguments from the Python side. Example:
    `v = ROOT.std.vector('int')()`
    `t.Branch('my_vector_branch', v)`
    """

    filename  = 'treebranch.root'
    treename  = 'mytree'
    arraysize = 10
    ival      = 7
    fval      = 7.

    # Setup
    @classmethod
    def setUpClass(cls):
        ROOT.gInterpreter.Declare("""
        struct MyStruct {
            int myint1;
            int myint2;
        };
        """)

    # Helpers
    def create_file_and_tree(self):
        f = ROOT.TFile(self.filename, 'RECREATE')
        t = ROOT.TTree(self.treename, self.treename)
        # Prevent double deletion of the tree (Python and C++ TFile)
        SetOwnership(t, False)

        return f,t

    def get_tree(self):
        f = ROOT.TFile(self.filename)
        t = f.Get(self.treename)
        SetOwnership(t, False)

        return f,t

    @staticmethod
    def fill_and_close(f, t):
        t.Fill()
        f.Write()
        f.Close()

    # Tests
    # Basic type and array do not actually need the pythonization,
    # but testing anyway for the sake of completeness
    def test01_write_basic_type_branch(self):
        f,t = self.create_file_and_tree()

        n = array('f', [ self.fval ])
        t.Branch('floatb', n, 'floatb/F')

        self.fill_and_close(f, t)

    def test02_read_basic_type_branch(self):
        f,t = self.get_tree()

        for entry in t:
            self.assertEqual(entry.floatb, self.fval)

    def test03_write_array_branch(self):
        f,t = self.create_file_and_tree()

        a = array('d', self.arraysize*[ self.fval ])
        t.Branch('arrayb', a, 'arrayb[' + str(self.arraysize) + ']/D')

        self.fill_and_close(f, t)

    def test04_read_array_branch(self):
        f,t = self.get_tree()

        for entry in t:
            for elem in entry.arrayb:
                self.assertEqual(elem, self.fval)

    def test05_write_numpy_array_branch(self):
        f,t = self.create_file_and_tree()

        a = np.array(self.arraysize*[ self.fval ]) # dtype='float64'
        t.Branch('arrayb', a, 'arrayb[' + str(self.arraysize) + ']/D')

        self.fill_and_close(f, t)

    def test06_read_numpy_array_branch(self):
        f,t = self.get_tree()

        for entry in t:
            for elem in entry.arrayb:
                self.assertEqual(elem, self.fval)

    # Struct and vector do benefit from the pythonization
    def test07_write_struct_single_branch(self):
        f,t = self.create_file_and_tree()

        ms = ROOT.MyStruct()
        ms.myint1, ms.myint2 = self.ival, 2*self.ival

        # Test overloads
        t.Branch('structb0', ms)
        t.Branch('structb1', ms, 32000)
        t.Branch('structb2', ms, 32000, 99)
        t.Branch('structb3', 'MyStruct', ms)
        t.Branch('structb4', 'MyStruct', ms, 32000)
        t.Branch('structb5', 'MyStruct', ms, 32000, 99)

        self.fill_and_close(f, t)

    def test08_read_struct_single_branch(self):
        f,t = self.get_tree()

        for entry in t:
            for ms in [ getattr(entry, 'structb' + str(i)) for i in range(6) ]:
                self.assertEqual(ms.myint1, self.ival)
                self.assertEqual(ms.myint2, 2*self.ival)

    def test09_write_struct_separate_branches(self):
        f,t = self.create_file_and_tree()

        ms = ROOT.MyStruct()
        ms.myint1, ms.myint2 = self.ival, 2*self.ival

        # Use `addressof` to get the address of the struct members
        t.Branch('myint1b', addressof(ms, 'myint1'), 'myint1balias/I')
        t.Branch('myint2b', addressof(ms, 'myint2'), 'myint2balias/I')

        self.fill_and_close(f, t)

    def test10_read_struct_separate_branches(self):
        f,t = self.get_tree()

        for entry in t:
            self.assertEqual(entry.myint1b, self.ival)
            self.assertEqual(entry.myint2b, 2*self.ival)
            # Test aliases
            self.assertEqual(entry.myint1balias, self.ival)
            self.assertEqual(entry.myint2balias, 2*self.ival)

    def test11_write_vector_branch(self):
        f,t = self.create_file_and_tree()

        v = ROOT.std.vector('double')(self.arraysize*[ self.fval ])

        # Test overloads
        t.Branch('vectorb0', v)
        t.Branch('vectorb1', v, 32000)
        t.Branch('vectorb2', v, 32000, 99)
        t.Branch('vectorb3', 'std::vector<double>', v)
        t.Branch('vectorb4', 'std::vector<double>', v, 32000)
        t.Branch('vectorb5', 'std::vector<double>', v, 32000, 99)

        self.fill_and_close(f, t)

    def test12_read_vector_branch(self):
        f,t = self.get_tree()

        for entry in t:
            for v in [ getattr(entry, 'vectorb' + str(i)) for i in range(6) ]:
                for elem in v:
                    self.assertEqual(elem, self.fval)

    def test13_write_fallback_case(self):
        f,t = self.create_file_and_tree()

        # Test an overload that uses the original Branch proxy
        l = ROOT.TList()
        s = ROOT.TString('one:two')
        a = s.Tokenize(ROOT.TString(':'))
        a.SetName('myobjarray')
        l.Add(a)
        t.Branch(l)

        self.fill_and_close(f, t)

    def test14_read_fallback_case(self):
        f,t = self.get_tree()

        for entry in t:
            self.assertEqual(entry.myobjarray_one.GetName(), 'one')
            self.assertEqual(entry.myobjarray_two.GetName(), 'two')

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


<commit_before>#!/usr/bin/env python3

"""
BehaviorModule
^^^^^^^^^^^^^^
.. moduleauthor:: Martin Poppinga <1popping@informatik.uni-hamburg.de>

Starts the body behavior
"""

import actionlib
import os
import rospy
from bitbots_msgs.msg import DynUpAction
from tf2_geometry_msgs import PoseStamped
from humanoid_league_msgs.msg import GameState, HeadMode, Strategy, TeamData,\
    RobotControlState, PoseWithCertainty, PoseWithCertaintyArray
from move_base_msgs.msg import MoveBaseActionFeedback, MoveBaseActionResult
from actionlib_msgs.msg import GoalID
from std_msgs.msg import Bool
from visualization_msgs.msg import Marker

from bitbots_blackboard.blackboard import BodyBlackboard
from bitbots_blackboard.async_service import AsyncServiceProxy
from dynamic_stack_decider import dsd
from geometry_msgs.msg import PoseWithCovarianceStamped, TwistWithCovarianceStamped
from bitbots_ros_patches.rate import Rate
from sensor_msgs.msg import PointCloud2
from std_msgs.msg import Float32
from nav_msgs.srv import GetPlan

if __name__ == "__main__":
    rospy.init_node("Bodybehavior")
    D = dsd.DSD(BodyBlackboard(), 'debug/dsd/body_behavior')

    D.blackboard.team_data.strategy_sender = rospy.Publisher("strategy", Strategy, queue_size=2)
    D.blackboard.team_data.time_to_ball_publisher = rospy.Publisher("time_to_ball", Float32, queue_size=2)
    D.blackboard.blackboard.head_pub = rospy.Publisher("head_mode", HeadMode, queue_size=10)
    D.blackboard.pathfinding.pathfinding_pub = rospy.Publisher('move_base_simple/goal', PoseStamped, queue_size=1)
    D.blackboard.pathfinding.pathfinding_cancel_pub = rospy.Publisher('move_base/cancel', GoalID, queue_size=1)
    D.blackboard.pathfinding.ball_obstacle_active_pub = rospy.Publisher("ball_obstacle_active", Bool, queue_size=1)
    D.blackboard.pathfinding.approach_marker_pub = rospy.Publisher("debug/approach_point", Marker, queue_size=10)
    D.blackboard.pathfinding.get_plan_service = AsyncServiceProxy("move_base/NavfnROS/make_plan", GetPlan)
    while not rospy.is_shutdown():
        try:
            D.blackboard.pathfinding.get_plan_service.service_proxy.wait_for_service(2.0)
        except rospy.ROSException as ex:
            rospy.logwarn("waiting for 'move_base/NavfnROS/make_plan' Service to become available...")
    D.blackboard.pathfinding.get_plan_service.callback = D.blackboard.pathfinding.path_to_ball_cb
    dirname = os.path.dirname(os.path.realpath(__file__))

    D.register_actions(os.path.join(dirname, "actions"))
    D.register_decisions(os.path.join(dirname, "decisions"))

    D.load_behavior(os.path.join(dirname, "main.dsd"))

    D.blackboard.dynup_action_client = actionlib.SimpleActionClient('dynup', DynUpAction)

    # TODO: callbacks away from the blackboard!
    rospy.Subscriber("balls_relative", PoseWithCertaintyArray, D.blackboard.world_model.balls_callback)
    rospy.Subscriber("ball_position_relative_filtered", PoseWithCovarianceStamped, D.blackboard.world_model.ball_filtered_callback)
    rospy.Subscriber("goal_posts_relative", PoseWithCertaintyArray, D.blackboard.world_model.goalposts_callback)
    rospy.Subscriber("gamestate", GameState, D.blackboard.gamestate.gamestate_callback)
    rospy.Subscriber("team_data", TeamData, D.blackboard.team_data.team_data_callback)
    rospy.Subscriber("pose_with_covariance", PoseWithCovarianceStamped, D.blackboard.world_model.pose_callback)
    rospy.Subscriber("robot_obstacles", PointCloud2, D.blackboard.world_model.robot_obstacle_callback)
    rospy.Subscriber("robot_state", RobotControlState, D.blackboard.blackboard.robot_state_callback)
    rospy.Subscriber(
        rospy.get_param("behavior/body/ball_movement_subscribe_topic"),
        TwistWithCovarianceStamped,
        D.blackboard.world_model.ball_twist_callback)
    rospy.Subscriber("move_base/feedback", MoveBaseActionFeedback, D.blackboard.pathfinding.feedback_callback)
    rospy.Subscriber("move_base/result", MoveBaseActionResult, D.blackboard.pathfinding.status_callback)

    rate = Rate(125)
    counter = 0
    while not rospy.is_shutdown():
        D.update()
        D.blackboard.team_data.publish_strategy()
        D.blackboard.team_data.publish_time_to_ball()
        counter = (counter + 1) % 125
        if counter == 0:
            D.blackboard.pathfinding.get_new_path_to_ball()
        rate.sleep()
<commit_msg>check gamestate playing before requesting path to ball<commit_after>#!/usr/bin/env python3

"""
BehaviorModule
^^^^^^^^^^^^^^
.. moduleauthor:: Martin Poppinga <1popping@informatik.uni-hamburg.de>

Starts the body behavior
"""

import actionlib
import os
import rospy
from bitbots_msgs.msg import DynUpAction
from tf2_geometry_msgs import PoseStamped
from humanoid_league_msgs.msg import GameState, HeadMode, Strategy, TeamData,\
    RobotControlState, PoseWithCertainty, PoseWithCertaintyArray
from move_base_msgs.msg import MoveBaseActionFeedback, MoveBaseActionResult
from actionlib_msgs.msg import GoalID
from std_msgs.msg import Bool
from visualization_msgs.msg import Marker

from bitbots_blackboard.blackboard import BodyBlackboard
from bitbots_blackboard.async_service import AsyncServiceProxy
from dynamic_stack_decider import dsd
from geometry_msgs.msg import PoseWithCovarianceStamped, TwistWithCovarianceStamped
from bitbots_ros_patches.rate import Rate
from sensor_msgs.msg import PointCloud2
from std_msgs.msg import Float32
from nav_msgs.srv import GetPlan

if __name__ == "__main__":
    rospy.init_node("Bodybehavior")
    D = dsd.DSD(BodyBlackboard(), 'debug/dsd/body_behavior')

    D.blackboard.team_data.strategy_sender = rospy.Publisher("strategy", Strategy, queue_size=2)
    D.blackboard.team_data.time_to_ball_publisher = rospy.Publisher("time_to_ball", Float32, queue_size=2)
    D.blackboard.blackboard.head_pub = rospy.Publisher("head_mode", HeadMode, queue_size=10)
    D.blackboard.pathfinding.pathfinding_pub = rospy.Publisher('move_base_simple/goal', PoseStamped, queue_size=1)
    D.blackboard.pathfinding.pathfinding_cancel_pub = rospy.Publisher('move_base/cancel', GoalID, queue_size=1)
    D.blackboard.pathfinding.ball_obstacle_active_pub = rospy.Publisher("ball_obstacle_active", Bool, queue_size=1)
    D.blackboard.pathfinding.approach_marker_pub = rospy.Publisher("debug/approach_point", Marker, queue_size=10)
    D.blackboard.pathfinding.get_plan_service = AsyncServiceProxy("move_base/NavfnROS/make_plan", GetPlan)
    while not rospy.is_shutdown():
        try:
            D.blackboard.pathfinding.get_plan_service.service_proxy.wait_for_service(2.0)
        except rospy.ROSException as ex:
            rospy.logwarn("waiting for 'move_base/NavfnROS/make_plan' Service to become available...")
    D.blackboard.pathfinding.get_plan_service.callback = D.blackboard.pathfinding.path_to_ball_cb
    dirname = os.path.dirname(os.path.realpath(__file__))

    D.register_actions(os.path.join(dirname, "actions"))
    D.register_decisions(os.path.join(dirname, "decisions"))

    D.load_behavior(os.path.join(dirname, "main.dsd"))

    D.blackboard.dynup_action_client = actionlib.SimpleActionClient('dynup', DynUpAction)

    # TODO: callbacks away from the blackboard!
    rospy.Subscriber("balls_relative", PoseWithCertaintyArray, D.blackboard.world_model.balls_callback)
    rospy.Subscriber("ball_position_relative_filtered", PoseWithCovarianceStamped, D.blackboard.world_model.ball_filtered_callback)
    rospy.Subscriber("goal_posts_relative", PoseWithCertaintyArray, D.blackboard.world_model.goalposts_callback)
    rospy.Subscriber("gamestate", GameState, D.blackboard.gamestate.gamestate_callback)
    rospy.Subscriber("team_data", TeamData, D.blackboard.team_data.team_data_callback)
    rospy.Subscriber("pose_with_covariance", PoseWithCovarianceStamped, D.blackboard.world_model.pose_callback)
    rospy.Subscriber("robot_obstacles", PointCloud2, D.blackboard.world_model.robot_obstacle_callback)
    rospy.Subscriber("robot_state", RobotControlState, D.blackboard.blackboard.robot_state_callback)
    rospy.Subscriber(
        rospy.get_param("behavior/body/ball_movement_subscribe_topic"),
        TwistWithCovarianceStamped,
        D.blackboard.world_model.ball_twist_callback)
    rospy.Subscriber("move_base/feedback", MoveBaseActionFeedback, D.blackboard.pathfinding.feedback_callback)
    rospy.Subscriber("move_base/result", MoveBaseActionResult, D.blackboard.pathfinding.status_callback)

    rate = Rate(125)
    counter = 0
    while not rospy.is_shutdown():
        D.update()
        D.blackboard.team_data.publish_strategy()
        D.blackboard.team_data.publish_time_to_ball()
        counter = (counter + 1) % 125
        if counter == 0 and D.blackboard.gamestate.is_game_state_equals(GameState.GAMESTATE_PLAYING):
            D.blackboard.pathfinding.get_new_path_to_ball()
        rate.sleep()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright (c) 2013-2014 Simon Jagoe
# All rights reserved.
#
# This software may be modified and distributed under the terms
# of the 3-clause BSD license.  See the LICENSE.txt file for details.
from __future__ import absolute_import, unicode_literals

from datetime import datetime
from enum import Enum
import sys
import time
import traceback

from .testing import unittest

STDOUT_LINE = '\nStdout:\n%s'
STDERR_LINE = '\nStderr:\n%s'

def _is_relevant_tb_level(tb):
    return '__unittest' in tb.tb_frame.f_globals

def _count_relevant_tb_levels(tb):
    length = 0
    while tb and not _is_relevant_tb_level(tb):
        length += 1
        tb = tb.tb_next
    return length

def _format_exception(err, is_failure, stdout=None, stderr=None):
    """Converts a sys.exc_info()-style tuple of values into a string."""
    exctype, value, tb = err
    # Skip test runner traceback levels
    while tb and _is_relevant_tb_level(tb):
        tb = tb.tb_next

    if is_failure:
        # Skip assert*() traceback levels
        length = _count_relevant_tb_levels(tb)
        msgLines = traceback.format_exception(exctype, value, tb, length)
    else:
        msgLines = traceback.format_exception(exctype, value, tb)

    if stdout:
        if not stdout.endswith('\n'):
            stdout += '\n'
        msgLines.append(STDOUT_LINE % stdout)
    if stderr:
        if not stderr.endswith('\n'):
            stderr += '\n'
        msgLines.append(STDERR_LINE % stderr)
    return ''.join(msgLines)

class TestCompletionStatus(Enum):
    success = 1
    failure = 2
    error = 3
    unexpected_success = 4
    expected_failure = 5
    skipped = 6

_successful_results = set([TestCompletionStatus.success,
                           TestCompletionStatus.expected_failure,
                           TestCompletionStatus.skipped])

class TestResult(object):

    def __init__(self, test_class, test_method_name, status,  # started_time,
                 completed_time, exception=None, message=None):
        self.test_class = test_class
        self.test_method_name = test_method_name
        self.status = status
        self.exception = exception
        self.message = message
        # self.started_time = started_time
        self.completed_time = completed_time

    @classmethod
    def from_test_case(cls, test_case, status,  # started_time,
                       exception=None,
                       message=None, stdout=None, stderr=None):
        test_class = type(test_case)
        test_method_name = test_case._testMethodName
        if exception is not None:
            is_failure = exception is test_case.failureException
            exception = _format_exception(
                exception, is_failure, stdout, stderr)
        completed_time = datetime.utcnow()
        return cls(test_class, test_method_name, status,  # started_time,
                   completed_time, exception, message)

    def __repr__(self):
        return '<{0} class={1}, method={2}>'.format(
            type(self).__name__, self.test_class.__name__,
            self.test_method_name)

    @property
    def test(self):
        return self.test_class(self.test_method_name)

    def get_test_description(self, show_doc_first_line):
        test = self.test
        doc_first_line = test.shortDescription()
        if show_doc_first_line and doc_first_line:
            return '\n'.join((str(test), doc_first_line))
        else:
            return str(test)

    def to_dict(self):
        return {
            'test_class': self.test_class,
            'test_method_name': self.test_method_name,
            'status': self.status,
            'exception': self.exception,
            'message': self.message,
            # 'started_time': self.started_time,
            'completed_time': self.completed_time,
        }

    @classmethod
    def from_dict(cls, data):
        return cls(**data)

# Temporary compatibility with unittest's runner
separator2 = '-' * 70

class ResultCollecter(object):

    # Temporary compatibility with unittest's runner
    separator2 = separator2

    def __init__(self, buffer=False):
        self.buffer = buffer
        self._handlers = []
        self.testsRun = 0
        self.expectedFailures = []
        self.unexpectedSuccesses = []
        self.skipped = []
        self.failures = []
        self.errors = []
        self._successful = True
        self.shouldStop = False

    def printErrors(self):
        pass

    def add_result_handler(self, handler):
        self._handlers.append(handler)

    def startTest(self, test):
        self.testsRun += 1
        for handler in self._handlers:
            handler.start_test(test)

    def stopTest(self, test):
        for handler in self._handlers:
            handler.stop_test(test)

    def startTestRun(self):
        for handler in self._handlers:
            handler.start_test_run()

    def stopTestRun(self):
        for handler in self._handlers:
            handler.stop_test_run()

    def _handle_result(self, test, status, exception=None, message=None):
        result = TestResult.from_test_case(
            test,
            status,
            exception=exception,
            message=message,
        )
        for handler in self._handlers:
            handler(result)
        if self._successful and result.status not in _successful_results:
            self._successful = False
        return result

    def addError(self, test, exception):
        result = self._handle_result(
            test, TestCompletionStatus.error, exception=exception)
        self.errors.append(result)

    def addFailure(self, test, exception):
        result = self._handle_result(
            test, TestCompletionStatus.failure, exception=exception)
        self.failures.append(result)

    def addSuccess(self, test):
        self._handle_result(test, TestCompletionStatus.success)

    def addSkip(self, test, reason):
        result = self._handle_result(
            test, TestCompletionStatus.skipped, message=reason)
        self.skipped.append(result)

    def addExpectedFailure(self, test, exception):
        result = self._handle_result(
            test, TestCompletionStatus.expected_failure, exception=exception)
        self.expectedFailures.append(result)

    def addUnexpectedSuccess(self, test):
        result = self._handle_result(
            test, TestCompletionStatus.unexpected_success)
        self.unexpectedSuccesses.append(result)

    def wasSuccessful(self):
        return self._successful

    def stop(self):
        self.shouldStop = True

class QuietTestResultHandler(object):
    separator1 = '=' * 70
    separator2 = separator2

    def __init__(self, test_count):
        self._test_count = test_count
        self.tests_run = 0
        self.descriptions = True
        self.expectedFailures = expectedFailures = []
        self.unexpectedSuccesses = unexpectedSuccesses = []
        self.skipped = skipped = []
        self.failures = failures = []
        self.errors = errors = []
        self._collectors = {
            TestCompletionStatus.failure: failures,
            TestCompletionStatus.error: errors,
            TestCompletionStatus.unexpected_success: unexpectedSuccesses,
            TestCompletionStatus.expected_failure: expectedFailures,
            TestCompletionStatus.skipped: skipped,
        }

    def start_test(self, test):
        self.tests_run += 1

    def stop_test(self, test):
        pass

    def start_test_run(self):
        pass

    def stop_test_run(self):
        self.print_errors()

    def print_errors(self):
        self.stream.writeln()
        self.print_error_list('ERROR', self.errors)
        self.print_error_list('FAIL', self.failures)

    def print_error_list(self, error_kind, errors):
        for result in errors:
            self.stream.writeln(self.separator1)
            self.stream.writeln(
                '%s: %s' % (error_kind, result.get_test_description(
                    self.descriptions)))
            self.stream.writeln(self.separator2)
            self.stream.writeln(result.exception)

    def __call__(self, result):
        collector = self._collectors.get(result.status)
        if collector is not None:
            collector.append(result)

class StandardTestResultHandler(QuietTestResultHandler):

    _result_formats = {
        TestCompletionStatus.success: '.',
        TestCompletionStatus.failure: 'F',
        TestCompletionStatus.error: 'E',
        TestCompletionStatus.unexpected_success: 'u',
        TestCompletionStatus.expected_failure: 'x',
        TestCompletionStatus.skipped: 's',
    }

    def __init__(self, test_count):
        super(StandardTestResultHandler, self).__init__(test_count)
        from unittest.runner import _WritelnDecorator
        self.stream = _WritelnDecorator(sys.stderr)

    def stop_test_run(self):
        self.stream.write('\n')
        super(StandardTestResultHandler, self).stop_test_run()

    def __call__(self, result):
        super(StandardTestResultHandler, self).__call__(result)
        self.stream.write(self._result_formats[result.status])
        self.stream.flush()

class TextTestResult(unittest.TextTestResult):
    """A simple extension to ``unittest.TextTestResult`` that displays
    progression of testing when run in verbose mode.

    """

    def __init__(self, total_tests, stream, descriptions, verbosity):
        """Create a TextTestResult. The parameters ``stream``, ``descriptions``
        and ``verbosity`` are as in ``unittest.TextTestResult``.

        Parameters
        ----------
        total_tests : int
            The total number of tests in the suite to be run, as
            returned by ``Suite.countTestCases()``

        """
        self._total_tests = total_tests
        super(TextTestResult, self).__init__(stream, descriptions, verbosity)

    def startTest(self, test):
        if self.showAll:
            padding = len(str(self._total_tests))
            prefix = '[{timestamp}] ({run: >{padding}d}/{total:d}) '.format(
                timestamp=time.ctime(),
                run=self.testsRun + 1,
                padding=padding,
                total=self._total_tests,
            )
            self.stream.write(prefix)
        super(TextTestResult, self).startTest(test)

    startTest.__doc__ = unittest.TextTestResult.startTest.__doc__
<commit_msg>Move get_test_description to result handler<commit_after># -*- coding: utf-8 -*-
# Copyright (c) 2013-2014 Simon Jagoe
# All rights reserved.
#
# This software may be modified and distributed under the terms
# of the 3-clause BSD license.  See the LICENSE.txt file for details.
from __future__ import absolute_import, unicode_literals

from datetime import datetime
from enum import Enum
import sys
import time
import traceback

from .testing import unittest

STDOUT_LINE = '\nStdout:\n%s'
STDERR_LINE = '\nStderr:\n%s'

def _is_relevant_tb_level(tb):
    return '__unittest' in tb.tb_frame.f_globals

def _count_relevant_tb_levels(tb):
    length = 0
    while tb and not _is_relevant_tb_level(tb):
        length += 1
        tb = tb.tb_next
    return length

def _format_exception(err, is_failure, stdout=None, stderr=None):
    """Converts a sys.exc_info()-style tuple of values into a string."""
    exctype, value, tb = err
    # Skip test runner traceback levels
    while tb and _is_relevant_tb_level(tb):
        tb = tb.tb_next

    if is_failure:
        # Skip assert*() traceback levels
        length = _count_relevant_tb_levels(tb)
        msgLines = traceback.format_exception(exctype, value, tb, length)
    else:
        msgLines = traceback.format_exception(exctype, value, tb)

    if stdout:
        if not stdout.endswith('\n'):
            stdout += '\n'
        msgLines.append(STDOUT_LINE % stdout)
    if stderr:
        if not stderr.endswith('\n'):
            stderr += '\n'
        msgLines.append(STDERR_LINE % stderr)
    return ''.join(msgLines)

class TestCompletionStatus(Enum):
    success = 1
    failure = 2
    error = 3
    unexpected_success = 4
    expected_failure = 5
    skipped = 6

_successful_results = set([TestCompletionStatus.success,
                           TestCompletionStatus.expected_failure,
                           TestCompletionStatus.skipped])

class TestResult(object):

    def __init__(self, test_class, test_method_name, status,  # started_time,
                 completed_time, exception=None, message=None):
        self.test_class = test_class
        self.test_method_name = test_method_name
        self.status = status
        self.exception = exception
        self.message = message
        # self.started_time = started_time
        self.completed_time = completed_time

    @classmethod
    def from_test_case(cls, test_case, status,  # started_time,
                       exception=None,
                       message=None, stdout=None, stderr=None):
        test_class = type(test_case)
        test_method_name = test_case._testMethodName
        if exception is not None:
            is_failure = exception is test_case.failureException
            exception = _format_exception(
                exception, is_failure, stdout, stderr)
        completed_time = datetime.utcnow()
        return cls(test_class, test_method_name, status,  # started_time,
                   completed_time, exception, message)

    def __repr__(self):
        return '<{0} class={1}, method={2}>'.format(
            type(self).__name__, self.test_class.__name__,
            self.test_method_name)

    @property
    def test(self):
        return self.test_class(self.test_method_name)

    def to_dict(self):
        return {
            'test_class': self.test_class,
            'test_method_name': self.test_method_name,
            'status': self.status,
            'exception': self.exception,
            'message': self.message,
            # 'started_time': self.started_time,
            'completed_time': self.completed_time,
        }

    @classmethod
    def from_dict(cls, data):
        return cls(**data)

# Temporary compatibility with unittest's runner
separator2 = '-' * 70

class ResultCollecter(object):

    # Temporary compatibility with unittest's runner
    separator2 = separator2

    def __init__(self, buffer=False):
        self.buffer = buffer
        self._handlers = []
        self.testsRun = 0
        self.expectedFailures = []
        self.unexpectedSuccesses = []
        self.skipped = []
        self.failures = []
        self.errors = []
        self._successful = True
        self.shouldStop = False

    def printErrors(self):
        pass

    def add_result_handler(self, handler):
        self._handlers.append(handler)

    def startTest(self, test):
        self.testsRun += 1
        for handler in self._handlers:
            handler.start_test(test)

    def stopTest(self, test):
        for handler in self._handlers:
            handler.stop_test(test)

    def startTestRun(self):
        for handler in self._handlers:
            handler.start_test_run()

    def stopTestRun(self):
        for handler in self._handlers:
            handler.stop_test_run()

    def _handle_result(self, test, status, exception=None, message=None):
        result = TestResult.from_test_case(
            test,
            status,
            exception=exception,
            message=message,
        )
        for handler in self._handlers:
            handler(result)
        if self._successful and result.status not in _successful_results:
            self._successful = False
        return result

    def addError(self, test, exception):
        result = self._handle_result(
            test, TestCompletionStatus.error, exception=exception)
        self.errors.append(result)

    def addFailure(self, test, exception):
        result = self._handle_result(
            test, TestCompletionStatus.failure, exception=exception)
        self.failures.append(result)

    def addSuccess(self, test):
        self._handle_result(test, TestCompletionStatus.success)

    def addSkip(self, test, reason):
        result = self._handle_result(
            test, TestCompletionStatus.skipped, message=reason)
        self.skipped.append(result)

    def addExpectedFailure(self, test, exception):
        result = self._handle_result(
            test, TestCompletionStatus.expected_failure, exception=exception)
        self.expectedFailures.append(result)

    def addUnexpectedSuccess(self, test):
        result = self._handle_result(
            test, TestCompletionStatus.unexpected_success)
        self.unexpectedSuccesses.append(result)

    def wasSuccessful(self):
        return self._successful

    def stop(self):
        self.shouldStop = True

class QuietTestResultHandler(object):
    separator1 = '=' * 70
    separator2 = separator2

    def __init__(self, test_count):
        self._test_count = test_count
        self.tests_run = 0
        self.descriptions = True
        self.expectedFailures = expectedFailures = []
        self.unexpectedSuccesses = unexpectedSuccesses = []
        self.skipped = skipped = []
        self.failures = failures = []
        self.errors = errors = []
        self._collectors = {
            TestCompletionStatus.failure: failures,
            TestCompletionStatus.error: errors,
            TestCompletionStatus.unexpected_success: unexpectedSuccesses,
            TestCompletionStatus.expected_failure: expectedFailures,
            TestCompletionStatus.skipped: skipped,
        }

    def get_test_description(self, test):
        doc_first_line = test.shortDescription()
        if self.descriptions and doc_first_line:
            return '\n'.join((str(test), doc_first_line))
        else:
            return str(test)

    def start_test(self, test):
        self.tests_run += 1

    def stop_test(self, test):
        pass

    def start_test_run(self):
        pass

    def stop_test_run(self):
        self.print_errors()

    def print_errors(self):
        self.stream.writeln()
        self.print_error_list('ERROR', self.errors)
        self.print_error_list('FAIL', self.failures)

    def print_error_list(self, error_kind, errors):
        for result in errors:
            self.stream.writeln(self.separator1)
            self.stream.writeln(
                '%s: %s' % (error_kind, self.get_test_description(
                    result.test)))
            self.stream.writeln(self.separator2)
            self.stream.writeln(result.exception)

    def __call__(self, result):
        collector = self._collectors.get(result.status)
        if collector is not None:
            collector.append(result)

class StandardTestResultHandler(QuietTestResultHandler):

    _result_formats = {
        TestCompletionStatus.success: '.',
        TestCompletionStatus.failure: 'F',
        TestCompletionStatus.error: 'E',
        TestCompletionStatus.unexpected_success: 'u',
        TestCompletionStatus.expected_failure: 'x',
        TestCompletionStatus.skipped: 's',
    }

    def __init__(self, test_count):
        super(StandardTestResultHandler, self).__init__(test_count)
        from unittest.runner import _WritelnDecorator
        self.stream = _WritelnDecorator(sys.stderr)

    def stop_test_run(self):
        self.stream.write('\n')
        super(StandardTestResultHandler, self).stop_test_run()

    def __call__(self, result):
        super(StandardTestResultHandler, self).__call__(result)
        self.stream.write(self._result_formats[result.status])
        self.stream.flush()

class TextTestResult(unittest.TextTestResult):
    """A simple extension to ``unittest.TextTestResult`` that displays
    progression of testing when run in verbose mode.

    """

    def __init__(self, total_tests, stream, descriptions, verbosity):
        """Create a TextTestResult. The parameters ``stream``, ``descriptions``
        and ``verbosity`` are as in ``unittest.TextTestResult``.

        Parameters
        ----------
        total_tests : int
            The total number of tests in the suite to be run, as
            returned by ``Suite.countTestCases()``

        """
        self._total_tests = total_tests
        super(TextTestResult, self).__init__(stream, descriptions, verbosity)

    def startTest(self, test):
        if self.showAll:
            padding = len(str(self._total_tests))
            prefix = '[{timestamp}] ({run: >{padding}d}/{total:d}) '.format(
                timestamp=time.ctime(),
                run=self.testsRun + 1,
                padding=padding,
                total=self._total_tests,
            )
            self.stream.write(prefix)
        super(TextTestResult, self).startTest(test)

    startTest.__doc__ = unittest.TextTestResult.startTest.__doc__
<|endoftext|>


<commit_before>'''
Contains WDM versions of all models and frameworks
'''
from concentration import CMRelation
from halo_model import HaloModel
from hmf._cache import cached_property
from copy import deepcopy
from hmf.filters import TopHat
import numpy as np
from scipy import integrate as intg
from hmf.wdm import MassFunctionWDM
from hmf._framework import get_model
import sys

#===============================================================================
# C-M relations
#===============================================================================
def CMRelationWDMRescaled(name,**kwargs):
    """
    Class factory for Rescaled CM relations.
    """
    x = getattr(sys.modules["halomod.concentration"], name)

    # class CMRelationWDMRescale(x):
    #     """
    #     Base class for simply rescaled concentration-mass relations (cf. Schneider
    #     2013, Bose+15)
    #     """
    def __init__(self, m_hm, **kwargs):
        super(self.__class__, self).__init__(**kwargs)
        self.m_hm = m_hm

    def cm(self, m):
        cm = super(self.__class__, self).cm(m)
        g1 = self.params['g1']
        g2 = self.params['g2']
        b0 = self.params['beta0']
        b1 = self.params['beta1']
        return cm * (1 + g1 * self.m_hm / m) ** (-g2)*(1+self.z)**(b0*self.z-b1)

    K = type(name + "WDM", (x,),{})
    K._defaults.update({"g1":60, "g2":0.17,"beta0":0.026,"beta1":0.04})

    K.__init__ = __init__
    K.cm = cm
    return K(**kwargs)

#===============================================================================
# Framework
#===============================================================================
class HaloModelWDM(HaloModel, MassFunctionWDM):
    """
    This class is a derivative of HaloModel which sets a few defaults that make
    more sense for a WDM model, and also implements the framework to include a
    smooth component
    """

    def __init__(self, **kw):
        kw.setdefault("cm_relation", "DuffyWDM")
        super(HaloModelWDM, self).__init__(**kw)

    @cached_property("_wdm", "dlog10m")
    def M(self):
        """ Halo masses in M_sun/h """
        return 10 ** np.arange(self._wdm.m_fs, 18, self.dlog10m) * u.MsolMass / self._hunit

    @cached_property("matter_power")
    def p_lin(self):
        """ The Schneider12 version of P_Lin = P_halofit * W(kR)"""
        f = TopHat(None, None, None, None)
        return np.exp(self.matter_power) * f.k_space(np.exp(self.lnk) * 2.0)

    @cached_property("rho_gtm", "mean_dens", "_wdm")
    def f_halos(self):
        """
        The total fraction of mass bound up in halos
        """
        return self.rho_gtm[0] / self.mean_dens

    @cached_property("f_halos", "power_ss", "power_sh", "power_hh")
    def power_matter(self):
        return (1 - self.f_halos) ** 2 * self.power_ss + \
            2 * (1 - self.f_halos) * self.f_halos * self.power_sh + \
            self.f_halos ** 2 * self.power_hh

    @cached_property("power_m_1h", "power_m_2h")
    def power_hh(self):
        return self.power_m_1h + self.power_m_2h

    @cached_property("_wdm", "dndm", "M", "lnk", "profile", "dlog10m", "f_halos", "mean_dens")
    def power_m_1h(self):
        integrand = self.dndm_rescaled * self.M ** 3
        u = self.profile.u(self.k, self.M, norm="m") ** 2
        out = np.zeros_like(self.k)
        for i, k in enumerate(self.k):
            r = (np.pi / k)  # half the radius
            mmin = self.filter_mod.radius_to_mass(r)
            if np.any(self.M > mmin):
                integ = integrand[self.M > mmin] * u[i, self.M > mmin]
                out[i] = intg.simps(integ, dx=self.dlog10m) * np.log(10)
            else:
                out[i] = 0.0
        return out / (self.f_halos * self.mean_dens) ** 2

    @cached_property("dndm", "_wdm", "bias", "lnk", "M", "profile", "f_halos", "mean_dens",
                     "matter_power")
    def power_m_2h(self):
        # Only use schneider for now
        integrand = self.M ** 2 * self.dndm_rescaled * self.bias
        out = np.zeros_like(self.k)
        u = self.profile.u(np.exp(self.k), self.M, norm="m")
        for i, k in enumerate(self.k):
            integ = integrand * u[i, :]
            out[i] = (intg.simps(integ, dx=self.dlog10m) * np.log(10))
        return self.p_lin * out ** 2 / (self.f_halos * self.mean_dens) ** 2

    @cached_property("dndm", "_wdm", "lnk", "M", "profile", "f_halos", "mean_dens",
                     "bias", "bias_smooth")
    def power_sh(self):
        integrand = self.M ** 2 * self.dndm_rescaled * self.bias
        out = np.zeros_like(self.k)
        u = self.profile.u(self.k, c.M, norm="m")
        for i, k in enumerate(self.k):
            integ = integrand * u [i, :]
            out[i] = intg.simps(integ, dx=self.dlog10m) * np.log(10)

        return self.bias_smooth * self.p_lin * out / (self.f_halos * self.mean_dens)

    @cached_property("bias_smooth", "matter_power")
    def power_ss(self):
        return self.bias_smooth ** 2 * np.exp(self.nonlinear_power)  # self.p_lin

    @cached_property("dndm", "bias", "M", "_wdm", "mean_dens", "f_halos")
    def bias_smooth(self):
        integrand = self.M ** 2 * self.dndm_rescaled * self.bias
        integral = intg.simps(integrand, dx=self.dlog10m) * np.log(10)
        return (1 / (1 - self.f_halos)) * (1 - integral / self.mean_dens)

    @cached_property("_wdm")
    def cm(self):
        kwargs = dict(nu=self.nu, z=self.z, growth=self.growth_model,
                      M=self.M, **self.cm_params)
        if np.issubclass_(self.cm_relation,CMRelation):
            if self.cm_relation.__class__.__name__.endswith("WDM"):
                cm = self.cm_relation(m_hm=self._wdm.m_hm, **kwargs)
            else:
                cm = self.cm_relation(**kwargs)
        elif self.cm_relation.endswith("WDM"):
            cm = CMRelationWDMRescaled(self.cm_relation[:-3],m_hm=self._wdm.m_hm, **kwargs)
        else:
            cm = get_model(self.cm_relation, "halomod.concentration", **kwargs)

        return cm
<commit_msg>removed extraneous methods after refactoring<commit_after>'''
Contains WDM versions of all models and frameworks
'''
from concentration import CMRelation
from halo_model import HaloModel
from hmf._cache import cached_property
from copy import deepcopy
from hmf.filters import TopHat
import numpy as np
from scipy import integrate as intg
from hmf.wdm import MassFunctionWDM
from hmf._framework import get_model
import sys
import astropy.units as u

#===============================================================================
# C-M relations
#===============================================================================
def CMRelationWDMRescaled(name,**kwargs):
    """
    Class factory for Rescaled CM relations.
    """
    x = getattr(sys.modules["halomod.concentration"], name)

    # class CMRelationWDMRescale(x):
    #     """
    #     Base class for simply rescaled concentration-mass relations (cf. Schneider
    #     2013, Bose+15)
    #     """
    def __init__(self, m_hm, **kwargs):
        super(self.__class__, self).__init__(**kwargs)
        self.m_hm = m_hm

    def cm(self, m):
        cm = super(self.__class__, self).cm(m)
        g1 = self.params['g1']
        g2 = self.params['g2']
        b0 = self.params['beta0']
        b1 = self.params['beta1']
        return cm * (1 + g1 * self.m_hm / m) ** (-g2)*(1+self.z)**(b0*self.z-b1)

    K = type(name + "WDM", (x,),{})
    K._defaults.update({"g1":60, "g2":0.17,"beta0":0.026,"beta1":0.04})

    K.__init__ = __init__
    K.cm = cm
    return K(**kwargs)

#===============================================================================
# Framework
#===============================================================================
class HaloModelWDM(HaloModel, MassFunctionWDM):
    """
    This class is a derivative of HaloModel which sets a few defaults that make
    more sense for a WDM model, and also implements the framework to include a
    smooth component
    """

    def __init__(self, **kw):
        kw.setdefault("cm_relation", "DuffyWDM")
        super(HaloModelWDM, self).__init__(**kw)

    # @cached_property("_wdm", "dlog10m")
    # def M(self):
    #     """ Halo masses in M_sun/h """
    #     return 10 ** np.arange(np.log10(self._wdm.m_fs.value), 18, self.dlog10m) * u.MsolMass / self._hunit

    @cached_property("rho_gtm", "mean_dens", "_wdm")
    def f_halos(self):
        """
        The total fraction of mass bound up in halos
        """
        return self.rho_gtm[0] / self.mean_dens

    @cached_property("f_halos","power_ss","power_sh")
    def power_mm(self):
        power_hh = super(HaloModelWDM,self).power_mm
    # @cached_property("f_halos", "power_ss", "power_sh", "power_hh")
    # def power_matter(self):
        return (1 - self.f_halos) ** 2 * self.power_ss + \
            2 * (1 - self.f_halos) * self.f_halos * self.power_sh + \
            self.f_halos ** 2 * power_hh

    # @cached_property("power_m_1h", "power_m_2h")
    # def power_hh(self):
    #     return self.power_m_1h + self.power_m_2h
    #
    # @cached_property("_wdm", "dndm", "M", "lnk", "profile", "dlog10m", "f_halos", "mean_dens")
    # def power_m_1h(self):
    #     integrand = self.dndm_rescaled * self.M ** 3
    #     u = self.profile.u(self.k, self.M, norm="m") ** 2
    #     out = np.zeros_like(self.k)
    #     for i, k in enumerate(self.k):
    #         r = (np.pi / k)  # half the radius
    #         mmin = self.filter_mod.radius_to_mass(r)
    #         if np.any(self.M > mmin):
    #             integ = integrand[self.M > mmin] * u[i, self.M > mmin]
    #             out[i] = intg.simps(integ, dx=self.dlog10m) * np.log(10)
    #         else:
    #             out[i] = 0.0
    #     return out / (self.f_halos * self.mean_dens) ** 2
    #
    # @cached_property("dndm", "_wdm", "bias", "lnk", "M", "profile", "f_halos", "mean_dens",
    #                  "matter_power")
    # def power_m_2h(self):
    #     # Only use schneider for now
    #     integrand = self.M ** 2 * self.dndm_rescaled * self.bias
    #     out = np.zeros_like(self.k)
    #     u = self.profile.u(np.exp(self.k), self.M, norm="m")
    #     for i, k in enumerate(self.k):
    #         integ = integrand * u[i, :]
    #         out[i] = (intg.simps(integ, dx=self.dlog10m) * np.log(10))
    #     return self.p_lin * out ** 2 / (self.f_halos * self.mean_dens) ** 2

    @cached_property("M","dndm","bias","profile","k","dlog10m","bias_smooth","_power_halo_centres","mean_density")
    def power_sh(self):
        integrand = self.M ** 2 * self.dndm * self.bias*self.profile.u(self.k, self.M, norm="m")
        pch = intg.simps(integrand,dx=np.log(10)*self.dlog10m)
        return self.bias_smooth * self._power_halo_centres * pch / self.mean_density

    @cached_property("bias_smooth", "matter_power")
    def power_ss(self):
        return self.bias_smooth ** 2 * self._power_halo_centres

    @cached_property("f_halos", "bias_effective_matter")
    def bias_smooth(self):
        # integrand = self.M ** 2 * self.dndm * self.bias
        # integral = intg.simps(integrand, dx=self.dlog10m) * np.log(10)
        return (1-self.bias_effective_matter) / (1 - self.f_halos)

    @cached_property("_wdm")
    def cm(self):
        kwargs = dict(nu=self.nu, z=self.z, growth=self.growth_model,
                      M=self.M, **self.cm_params)
        if np.issubclass_(self.cm_relation,CMRelation):
            if self.cm_relation.__class__.__name__.endswith("WDM"):
                cm = self.cm_relation(m_hm=self._wdm.m_hm, **kwargs)
            else:
                cm = self.cm_relation(**kwargs)
        elif self.cm_relation.endswith("WDM"):
            cm = CMRelationWDMRescaled(self.cm_relation[:-3],m_hm=self._wdm.m_hm, **kwargs)
        else:
            cm = get_model(self.cm_relation, "halomod.concentration", **kwargs)

        return cm
<|endoftext|>


<commit_before>#!/usr/bin/python

'''
This is HashCollect! This tool will scrape md5 hashes out of a specified file or URL
@TekDefense
Ian Ahl | www.TekDefense.com | 1aN0rmus@tekDefense.com
'''

import httplib2, re, sys, argparse, urllib, urllib2

parser = argparse.ArgumentParser(description='Hash Collector')
parser.add_argument('-u', '--url', help='This option is used to search for hashes on a website')
parser.add_argument('-f', '--file', help='This option is used to import a file that contains hashes')
parser.add_argument('-o', '--output', help='This option will output the results to a file.')
args = parser.parse_args()

if args.url == None and args.file == None:
    parser.print_help()

if args.output:
    oFile = args.output
    print '[+] Printing results to file:', args.output
    o = open(oFile, "w")
    sys.stdout = o

if args.file:
    iFile = args.file
    fileImport =open(iFile)
    strFile=''
    for line in fileImport:
        strFile += line
    regexValue = re.compile('[a-fA-F0-9]{32}')
    regexSearch = re.findall(regexValue,strFile)
    for i in regexSearch:
        print i

if args.url:
    url = args.url
    h = httplib2.Http(".cache")
    resp, content = h.request((url), "GET")
    contentString = (str(content))
    regexValue = re.compile('[a-fA-F0-9]{32}')
    regexSearch = re.findall(regexValue,contentString)
    for i in regexSearch:
        print i

if __name__ == '__main__':
    pass
<commit_msg>Update hashCollect to allow custom regex<commit_after>#!/usr/bin/python

'''
This is HashCollect! This tool will scrape md5 hashes out of a specified file or URL
@TekDefense
Ian Ahl | www.TekDefense.com | 1aN0rmus@tekDefense.com
'''

import httplib2, re, sys, argparse, urllib, urllib2

parser = argparse.ArgumentParser(description='Hash Collector')
parser.add_argument('-u', '--url', help='This option is used to search for hashes on a website')
parser.add_argument('-f', '--file', help='This option is used to import a file that contains hashes')
parser.add_argument('-o', '--output', help='This option will output the results to a file.')
parser.add_argument('-r', '--regex', help='This option allows the user to set a custom regex value. Must incase in single or double quotes.')

args = parser.parse_args()
regVal = '[a-fA-F0-9]{32}'
listResults = []

if args.regex:
    regVal = str(args.regex)

if args.url == None and args.file == None:
    parser.print_help()

if args.output:
    oFile = args.output
    print '[+] Printing results to file:', args.output
    o = open(oFile, "w")
    sys.stdout = o

if args.file:
    iFile = args.file
    fileImport =open(iFile)
    strFile=''
    for line in fileImport:
        strFile += line
    regexValue = re.compile(regVal)
    regexSearch = re.findall(regexValue,strFile)
    for i in regexSearch:
        listResults.append(i)

if args.url:
    url = args.url
    h = httplib2.Http(".cache")
    resp, content = h.request((url), "GET")
    contentString = (str(content))
    regexValue = re.compile(regVal)
    regexSearch = re.findall(regexValue,contentString)
    for i in regexSearch:
        listResults.append(i)

listResults = list(set(listResults))
for i in listResults:
    print i

if __name__ == '__main__':
    pass
<|endoftext|>


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

import sys
from multiprocessing import Process
import time
import signal
from collections import Counter
import setproctitle

from .mixins import RoutesMixin, AppEventsMixin
from . import autoreload
from .log import logger
from . import constants

class Haven(RoutesMixin, AppEventsMixin):
    enable = True
    name = constants.NAME
    processes = None
    debug = False
    got_first_request = False
    blueprints = None

    def __init__(self):
        RoutesMixin.__init__(self)
        AppEventsMixin.__init__(self)
        self.processes = []
        self.blueprints = list()

    def register_blueprint(self, blueprint):
        blueprint.register_to_app(self)

    def run(self, host=None, port=None, debug=None, use_reloader=None, workers=None):
        self._validate_cmds()

        if host is None:
            host = constants.SERVER_HOST
        if port is None:
            port = constants.SERVER_PORT
        if debug is not None:
            self.debug = debug

        use_reloader = use_reloader if use_reloader is not None else self.debug
        if use_reloader and workers is not None:
            # 当 use_reloader 打开的时候，workers会强制变为None
            logger.warning(
                'use_reloader is %s, workers will be changed from %s to None.',
                use_reloader, workers
            )
            workers = None

        def run_wrapper():
            logger.info('Running server on %s:%s, debug: %s, use_reloader: %s, workers: %s',
                        host, port, self.debug, use_reloader, workers)

            self._prepare_server(host, port)
            if workers is not None:
                setproctitle.setproctitle(self._make_proc_name('master'))
                # 只能在主线程里面设置signals
                self._handle_parent_proc_signals()
                self._fork_workers(workers)
            else:
                self._try_serve_forever(True)

        if use_reloader:
            autoreload.main(run_wrapper)
        else:
            run_wrapper()

    def acquire_got_first_request(self):
        pass

    def release_got_first_request(self):
        pass

    def repeat_timer(self, interval):
        raise NotImplementedError

    def _make_proc_name(self, subtitle):
        """
        获取进程名称
        :param subtitle:
        :return:
        """
        proc_name = '[%s %s:%s] %s' % (
            self.name,
            constants.NAME,
            subtitle,
            ' '.join([sys.executable] + sys.argv)
        )

        return proc_name

    def _validate_cmds(self):
        """
        确保 cmd 没有重复
        :return:
        """

        cmd_list = list(self.rule_map.keys())

        for bp in self.blueprints:
            cmd_list.extend(bp.rule_map.keys())

        duplicate_cmds = (Counter(cmd_list) - Counter(set(cmd_list))).keys()

        assert not duplicate_cmds, 'duplicate cmds: %s' % duplicate_cmds

    def _before_worker_run(self):
        self.events.create_worker()
        for bp in self.blueprints:
            bp.events.create_app_worker()

        self.events.repeat_timer()
        for bp in self.blueprints:
            bp.events.repeat_app_timer()

    def _try_serve_forever(self, main_process):
        # 无论是否有master，这里都是worker
        if not main_process:
            setproctitle.setproctitle(self._make_proc_name('worker'))
            self._handle_child_proc_signals()
        else:
            setproctitle.setproctitle(self._make_proc_name('main'))

        self._before_worker_run()

        try:
            self._serve_forever()
        except KeyboardInterrupt:
            pass
        except:
            logger.error('exc occur.', exc_info=True)

    def _fork_workers(self, workers):
        def start_worker_process():
            inner_p = Process(target=self._try_serve_forever, args=(False,))
            # 当前进程daemon默认是False，改成True将启动不了子进程
            # 但是子进程要设置daemon为True，这样父进程退出，子进程会被强制关闭
            # 现在父进程会在子进程之后推出，没必要设置了
            # inner_p.daemon = True
            inner_p.start()
            return inner_p

        for it in xrange(0, workers):
            p = start_worker_process()
            self.processes.append(p)

        while 1:
            for idx, p in enumerate(self.processes):
                if p and not p.is_alive():
                    self.processes[idx] = None

                    if self.enable:
                        p = start_worker_process()
                        self.processes[idx] = p

            if not filter(lambda x: x, self.processes):
                # 没活着的了
                break

            # 时间短点，退出的快一些
            time.sleep(0.1)

    def _handle_parent_proc_signals(self):
        def exit_handler(signum, frame):
            self.enable = False

            # 如果是终端直接CTRL-C，子进程自然会在父进程之后收到INT信号，不需要再写代码发送
            # 如果直接kill -INT $parent_pid，子进程不会自动收到INT
            # 所以这里可能会导致重复发送的问题，重复发送会导致一些子进程异常，所以在子进程内部有做重复处理判断。
            for p in self.processes:
                if p:
                    p.terminate()

        # INT, QUIT, TERM为强制结束
        signal.signal(signal.SIGINT, exit_handler)
        signal.signal(signal.SIGQUIT, exit_handler)
        signal.signal(signal.SIGTERM, exit_handler)

    def _handle_child_proc_signals(self):
        def exit_handler(signum, frame):
            # 防止重复处理KeyboardInterrupt，导致抛出异常
            if self.enable:
                self.enable = False
                raise KeyboardInterrupt

        # 强制结束，抛出异常终止程序进行
        signal.signal(signal.SIGINT, exit_handler)
        signal.signal(signal.SIGQUIT, exit_handler)
        signal.signal(signal.SIGTERM, exit_handler)

    def _prepare_server(self, host, port):
        raise NotImplementedError

    def _serve_forever(self):
        raise NotImplementedError
<commit_msg> 当 use_reloader 打开的时候，workers会强制变为None<commit_after># -*- coding: utf-8 -*-

import sys
from multiprocessing import Process
import time
import signal
from collections import Counter
import setproctitle

from .mixins import RoutesMixin, AppEventsMixin
from . import autoreload
from .log import logger
from . import constants

class Haven(RoutesMixin, AppEventsMixin):
    enable = True
    name = constants.NAME
    processes = None
    debug = False
    got_first_request = False
    blueprints = None

    def __init__(self):
        RoutesMixin.__init__(self)
        AppEventsMixin.__init__(self)
        self.processes = []
        self.blueprints = list()

    def register_blueprint(self, blueprint):
        blueprint.register_to_app(self)

    def run(self, host=None, port=None, debug=None, use_reloader=None, workers=None):
        self._validate_cmds()

        if host is None:
            host = constants.SERVER_HOST
        if port is None:
            port = constants.SERVER_PORT
        if debug is not None:
            self.debug = debug

        use_reloader = use_reloader if use_reloader is not None else self.debug
        if use_reloader and workers is not None:
            # 当 use_reloader 打开的时候，workers会强制变为None
            workers = None

        def run_wrapper():
            logger.info('Running server on %s:%s, debug: %s, use_reloader: %s, workers: %s',
                        host, port, self.debug, use_reloader, workers)

            self._prepare_server(host, port)
            if workers is not None:
                setproctitle.setproctitle(self._make_proc_name('master'))
                # 只能在主线程里面设置signals
                self._handle_parent_proc_signals()
                self._fork_workers(workers)
            else:
                self._try_serve_forever(True)

        if use_reloader:
            autoreload.main(run_wrapper)
        else:
            run_wrapper()

    def acquire_got_first_request(self):
        pass

    def release_got_first_request(self):
        pass

    def repeat_timer(self, interval):
        raise NotImplementedError

    def _make_proc_name(self, subtitle):
        """
        获取进程名称
        :param subtitle:
        :return:
        """
        proc_name = '[%s %s:%s] %s' % (
            self.name,
            constants.NAME,
            subtitle,
            ' '.join([sys.executable] + sys.argv)
        )

        return proc_name

    def _validate_cmds(self):
        """
        确保 cmd 没有重复
        :return:
        """

        cmd_list = list(self.rule_map.keys())

        for bp in self.blueprints:
            cmd_list.extend(bp.rule_map.keys())

        duplicate_cmds = (Counter(cmd_list) - Counter(set(cmd_list))).keys()

        assert not duplicate_cmds, 'duplicate cmds: %s' % duplicate_cmds

    def _before_worker_run(self):
        self.events.create_worker()
        for bp in self.blueprints:
            bp.events.create_app_worker()

        self.events.repeat_timer()
        for bp in self.blueprints:
            bp.events.repeat_app_timer()

    def _try_serve_forever(self, main_process):
        # 无论是否有master，这里都是worker
        if not main_process:
            setproctitle.setproctitle(self._make_proc_name('worker'))
            self._handle_child_proc_signals()
        else:
            setproctitle.setproctitle(self._make_proc_name('main'))

        self._before_worker_run()

        try:
            self._serve_forever()
        except KeyboardInterrupt:
            pass
        except:
            logger.error('exc occur.', exc_info=True)

    def _fork_workers(self, workers):
        def start_worker_process():
            inner_p = Process(target=self._try_serve_forever, args=(False,))
            # 当前进程daemon默认是False，改成True将启动不了子进程
            # 但是子进程要设置daemon为True，这样父进程退出，子进程会被强制关闭
            # 现在父进程会在子进程之后推出，没必要设置了
            # inner_p.daemon = True
            inner_p.start()
            return inner_p

        for it in xrange(0, workers):
            p = start_worker_process()
            self.processes.append(p)

        while 1:
            for idx, p in enumerate(self.processes):
                if p and not p.is_alive():
                    self.processes[idx] = None

                    if self.enable:
                        p = start_worker_process()
                        self.processes[idx] = p

            if not filter(lambda x: x, self.processes):
                # 没活着的了
                break

            # 时间短点，退出的快一些
            time.sleep(0.1)

    def _handle_parent_proc_signals(self):
        def exit_handler(signum, frame):
            self.enable = False

            # 如果是终端直接CTRL-C，子进程自然会在父进程之后收到INT信号，不需要再写代码发送
            # 如果直接kill -INT $parent_pid，子进程不会自动收到INT
            # 所以这里可能会导致重复发送的问题，重复发送会导致一些子进程异常，所以在子进程内部有做重复处理判断。
            for p in self.processes:
                if p:
                    p.terminate()

        # INT, QUIT, TERM为强制结束
        signal.signal(signal.SIGINT, exit_handler)
        signal.signal(signal.SIGQUIT, exit_handler)
        signal.signal(signal.SIGTERM, exit_handler)

    def _handle_child_proc_signals(self):
        def exit_handler(signum, frame):
            # 防止重复处理KeyboardInterrupt，导致抛出异常
            if self.enable:
                self.enable = False
                raise KeyboardInterrupt

        # 强制结束，抛出异常终止程序进行
        signal.signal(signal.SIGINT, exit_handler)
        signal.signal(signal.SIGQUIT, exit_handler)
        signal.signal(signal.SIGTERM, exit_handler)

    def _prepare_server(self, host, port):
        raise NotImplementedError

    def _serve_forever(self):
        raise NotImplementedError
<|endoftext|>


<commit_before>#!/usr/bin/env python

# An attempt to make a LCD API 1.0 compatible (Arduino) library for python/Raspberry Pi for displays
# based on a HD44780 compatible display using a PCF8574 compatible I2C controller. Arduino users use
# the LiquidCrystal_I2C library to support this, and well use that as a reference.
#
# http://playground.arduino.cc/Code/LCDAPI
# https://hmario.home.xs4all.nl/arduino/LiquidCrystal_I2C/

from time import sleep
import RPi.GPIO as GPIO
# TODO: which python I2C library to use?
# pyI2C - http://pyi2c.sourceforge.net/
# smbus - python-smbus package
# quick2wire.i2c - http://quick2wire.com/i2c-python/

# Time is in microseconds
# Wikipedia says max command execution time is 1.52ms
DEFAULT_CMD_DELAY  = 1550
DEFAULT_CHAR_DELAY = 50

# Taken from LiquidCrystal_I2C library and Wikipedia page:
# https://en.wikipedia.org/wiki/Hitachi_HD44780_LCD_controller
#
# Commands
LCD_CMD_CLEARDISPLAY = 0x01
LCD_CMD_CURSORHOME   = 0x02
LCD_CMD_ENTRYMODESET = 0x04
LCD_CMD_DISPLAYCONTROL = 0x08
LCD_CMD_CURSORSHIFT  = 0x10
LCD_CMD_FUNCTIONSET  = 0x20
LCD_CMD_SETCGRAMADDR = 0x40
LCD_CMD_SETDDRAMADDR = 0x80

# Entry Mode Flags
LCD_ENTRYRIGHT = 0x00
LCD_ENTRYLEFT  = 0x02
LCD_ENTRYSHIFTINCR = 0x01
LCD_ENTRYSHIFTDECR = 0x00

# Display Control Flags
LCD_DISPLAYON  = 0x04
LCD_DISPLAYOFF = 0x00
LCD_CURSORON   = 0x02
LCD_CURSOROFF  = 0x00
LCD_BLINKON    = 0x01
LCD_BLINKOFF   = 0x00

# Cursor Shift Flags
LCD_DISPLAYMOVE = 0x08
LCD_CURSORMOVE  = 0x00
LCD_SHIFTRIGHT  = 0x01
LCD_SHIFTLEFT   = 0x00

# Function Set Flags
LCD_8BITMODE = 0x10
LCD_4BITMODE = 0x00
LCD_2LINE = 0x08
LCD_1LINE = 0x00
LCD_5x10DOTS = 0x04
LCD_5x8DOTS  = 0x00

class hd44780_i2c():
  def _init_pi():
    # This make a big assumtion that GPIO.setmode() hasn't already been called, and
    # that someone really had their heart set on using mode GPIO.BOARD
    # RPi I2C pins are always on GPIO pins 2 (SDA) & 3 (SCL) using BCM mode
    GPIO.setmode(GPIO.BCM)
    GPIO.setup([2,3], GPIO.OUT, initial = GPIO.LOW)

  def _init_display():
    # Per http://web.stanford.edu/class/ee281/handouts/lcd_tutorial.pdf
    # Initialize the display to a known default state. Must send at least a function set
    # command, and optionally send entry mode, display control and clear display commands.
    # This code will:
    #   FUNCTION SET: 8-bit mode, 2 lines with 5x8 characters
    #   ENTRY MODE: L -> R with no display shift
    #   DISPLAY MODE: Turn display and cursor on, no cursor blink
    #
    # Execute the 'Initialize by Instruction' routine specified in the datasheet, since we won't assume
    # the system is supplying the necessary Vcc to trigger the power-on reset logic.  It'll also ensure
    # we're starting from a totally clean slate when we instantiate new instances of this class

    func_set = LCD_CMD_FUNCTIONSET | LCD_8BITMODE | LCD_2LINE | LCD_5x8DOTS
    entry_mode_set = LCD_CMD_ENTRYMODSET | LCD_ENTRYLEFT | LCD_ENTRYSHIFTDECR
    display_control_set = LCD_CMD_DISPLAYCONTROL | LCD_DISPLAYON | LCD_CURSORON | LCD_BLINKOFF

    # Wait at least 40ms after Vcc hits 2.7V
    sleep(0.1)
    self.command(0x30)
    sleep(0.01) # Default command delay too short
    self.command(0x30) # Per datasheet, default delay should be adequate here
    self.command(0x30)
    sleep(0.1) # for good measure

    # Reset instructions complete, now initialize
    self.command(func_set)
    self.command(entry_mode_set)
    self.command(display_control_set)
    self.clear()

  # Mandatory functions
  def __init__(self, i2c_addr, rows, cols, **kwargs):
    assert(i2c_addr > 0)
    assert(rows > 0)
    assert(cols > 0)

    # Note to self: a 20x4 display may be a 40x2 display
    self.i2c_addr = i2c_addr
    self.rows = rows
    self.cols = cols
    self.test = kwargs.get('test', False)

    if not self.test:
      # TODO: initialize GPIO pins, I2C library and display
      _init_pi()
      _init_display()

    self.set_delay()

  def set_delay(self, **kwargs):
    # Override the default library delays. kwargs can be one of 'cmd', or 'char'
    # to specify setting the delay for LCD commands, or sending characters.
    self.cmd_delay  = kwargs.get('cmd', DEFAULT_CMD_DELAY)
    self.char_delay = kwargs.get('char', DEFAULT_CHAR_DELAY)

# Clash with Python's print()?
#  def print(self, val):
#    # TODO: print the provided value on the display
#    pass

#  def println(self, val):
#    # Call print(), with trailing new line
#    self.print(str(val) + "\n")

  def write(self, val):
    # TODO: raw write value to the display
    pass

  def command(self, val):
    # TODO: Send command to display, for display-specific commands not covered in the API
    sleep(self.cmd_delay / 1000000) # Pause to ensure command is executed
    return bin(val)[2:].zfill(8)

  def clear(self):
    # clear display and return cursor to 0,0
    self.command(LCD_CMD_CLEARDISPLAY)

  def home(self):
    # set cursor position to 0,0 leaving display untouched
    self.command(LCD_CMD_CURSORHOME)

  def set_cursor(self, row, col):
    # move cursor to indicated position, row and col values falling outside the configured
    # display size will be set to 0 or the rows/colums value provided in the constructor
    if row < 0:
      row = 0

    if col < 0:
      col = 0

    if row > self.rows:
      row = self.rows

    if col > self.cols:
      col = self.cols

    if row == 0 and col == 0:
      self.home()
    else:
      # TODO: move cursor to the indicated position.  Should we send cursor home command and then
      # move, or determine current cursor position and move relative to that?
      pass

  def cursor_on(self):
    # set block cursor on
    self.command(LCD_CMD_DISPLAYCONTROL | LCD_CURSORON)

  def cursor_off(self):
    # set block cursor off
    self.command(LCD_CMD_DISPLAYCONTROL | LCD_CURSOROFF)

  def blink_on(self):
    # set blinking underline cursor on
    self.command(LCD_CMD_DISPLAYCONTROL | LCD_BLINKON)

  def blink_off(self):
    # set blinking underline cursor off
    self.command(LCD_CMD_DISPLAYCONTROL | LCD_BLINKOFF)

  # Optional functions
  def set_backlight(self, val):
    # TODO: set backlight brightnetss (0-255), where 0 = off
    pass

  def set_contrast(self, val):
    # TODO: set display contrast (0-255)
    pass

  def on(self):
    # turn display on and set backlight to ~ 75%
    self.command(LCD_CMD_DISPLAYCONTROL | LCD_DISPLAYON)
    self.set_backlight(192)

  def off(self):
    # set backlight to 0 and turn display off
    self.set_backlight(0)
    self.command(LCD_CMD_DISPLAYCONTROL | LCD_DISPLAYOFF)

  def status(self):
    # TODO: return status of the display. API docs say:
    #   Returns the FIFO buffer on the robot-electronics display
    #   Can be used to get r/w status of hardwired displays
    pass

  # Extended functions left unimplemented:
  #  load_custom_characters(self, char_num, rows)
  #  keypad(self)
  #  printstr(self)
  # These seem like a cool idea, and may implment in the future:
  #  init_bargraph(self, type)
  #  draw_horizontal_graph(row, col, len, end)
  #  draw_vertical_graph(row, col, len, end)

if __name__ == "__main__":
  cls = hd44780_i2c(0x37, 20, 4, test = True)
  print(cls.off())
  print(cls.blink_on())
<commit_msg>closer, but not totally right. Some control and characters still aren't quite right<commit_after>#!/usr/bin/env python

# An attempt to make a LCD API 1.0 compatible (Arduino) library for python/Raspberry Pi for displays
# based on a HD44780 compatible display using a PCF8574 compatible I2C controller. Arduino users use
# the LiquidCrystal_I2C library to support this, and well use that as a reference.
#
# http://playground.arduino.cc/Code/LCDAPI
# https://hmario.home.xs4all.nl/arduino/LiquidCrystal_I2C/

# 0x08 - backlight on
# 0x04 - E line high
# 0x02 - R/W line high
# 0x01 - RS line high

from time import sleep
import smbus

# Time is in microseconds
# Wikipedia says max command execution time is 1.52ms
DEFAULT_CMD_DELAY  = 1550
DEFAULT_CHAR_DELAY = 50

# Control status of the Register Select (RS) line
LCD_REG_CMD  = 0x00
LCD_REG_DATA = 0x01

# OR these values with LCD_CMD_SETDDRAMADDR to send the Set DDRAM Address command
LCD_LINE1_ADDR = 0x00
LCD_LINE2_ADDR = 0x40
LCD_LINE3_ADDR = 0x14 # 20 chars into line 1
LCD_LINE4_ADDR = 0x54 # 20 chars into line 2

# Taken from LiquidCrystal_I2C library and Wikipedia page:
# https://en.wikipedia.org/wiki/Hitachi_HD44780_LCD_controller
#
# Commands
LCD_CMD_CLEARDISPLAY = 0x01
LCD_CMD_CURSORHOME   = 0x02
LCD_CMD_ENTRYMODESET = 0x04
LCD_CMD_DISPLAYCONTROL = 0x08
LCD_CMD_CURSORSHIFT  = 0x10
LCD_CMD_FUNCTIONSET  = 0x20
LCD_CMD_SETCGRAMADDR = 0x40
LCD_CMD_SETDDRAMADDR = 0x80

# Entry Mode Flags
LCD_ENTRYRIGHT = 0x00
LCD_ENTRYLEFT  = 0x02
LCD_ENTRYSHIFTINCR = 0x01
LCD_ENTRYSHIFTDECR = 0x00

# Display Control Flags
LCD_DISPLAYON  = 0x04
LCD_DISPLAYOFF = 0x00
LCD_CURSORON   = 0x02
LCD_CURSOROFF  = 0x00
LCD_BLINKON    = 0x01
LCD_BLINKOFF   = 0x00

# Cursor Shift Flags
LCD_DISPLAYMOVE = 0x08
LCD_CURSORMOVE  = 0x00
LCD_SHIFTRIGHT  = 0x01
LCD_SHIFTLEFT   = 0x00

# Function Set Flags
LCD_8BITMODE = 0x10
LCD_4BITMODE = 0x00
LCD_2LINE = 0x08
LCD_1LINE = 0x00
LCD_5x10DOTS = 0x04
LCD_5x8DOTS  = 0x00

# flags for backlight control
# Sending 0x08 turns backlight on, any other value where the
# fourth bit is turned off will turn off the backlight
LCD_BACKLIGHT   = 0x08
LCD_NOBACKLIGHT = 0x00

class hd44780_i2c():
  def _init_display(self):
    # Per http://web.stanford.edu/class/ee281/handouts/lcd_tutorial.pdf
    # Initialize the display to a known default state. Must send at least a function set command,
    # and optionally send entry mode, display control and clear display commands. Since 3 of the 8
    # I2C outputs are used for control (RS, R/W and E lines), we need to use 4-bit mode for sending
    # data so we can also specify the necessary control signals on the upper 4 bits of the byte.
    #
    # This code will:
    #   FUNCTION SET: 4-bit mode, 2 lines with 5x8 characters
    #   ENTRY MODE: L -> R with no display shift
    #   DISPLAY MODE: Turn display and cursor on, no cursor blink
    #
    # Execute the 'Initialize by Instruction' routine specified in the datasheet, since we won't assume
    # the system is supplying the necessary Vcc to trigger the power-on reset logic.  It'll also ensure
    # we're starting from a totally clean slate when we instantiate new instances of this class

    func_set = LCD_CMD_FUNCTIONSET | LCD_4BITMODE | LCD_2LINE | LCD_5x8DOTS
    entry_mode_set = LCD_CMD_ENTRYMODESET | LCD_ENTRYLEFT | LCD_ENTRYSHIFTDECR
    display_control_set = LCD_CMD_DISPLAYCONTROL | LCD_DISPLAYON | LCD_CURSOROFF | LCD_BLINKOFF

    # Wait at least 40ms after Vcc hits 2.7V
    sleep(0.1)
    self.command(0x30)
    sleep(0.01) # Default command delay too short
    self.command(0x30) # Per datasheet, default delay should be adequate here
    self.command(0x30)
    sleep(0.1) # for good measure

    # Reset instructions complete, now initialize
    self.command(func_set)
    self.command(entry_mode_set)
    self.command(display_control_set)
    self.set_backlight(LCD_BACKLIGHT)
    self.clear()

  # Mandatory functions
  def __init__(self, i2c_bus, i2c_addr, rows, cols, **kwargs):
    assert(i2c_bus >= 0)
    assert(i2c_addr > 0)
    assert(rows > 0)
    assert(cols > 0)

    # Note to self: a 20x4 display may be a 40x2 display
    self.i2c_addr = i2c_addr
    self.rows = rows
    self.cols = cols
    self.test = kwargs.get('test', False)
    self.set_delay()
    self.backlight = LCD_NOBACKLIGHT

    if not self.test:
      # initialize I2C library and display, there is no special
      # RPi setup needed other than enabling I2C in the kernel
      self.bus = smbus.SMBus(i2c_bus)
      self._init_display()

  def set_delay(self, **kwargs):
    # Override the default library delays. kwargs can be one of 'cmd', or 'char'
    # to specify setting the delay for LCD commands, or sending characters.
    self.cmd_delay  = kwargs.get('cmd', DEFAULT_CMD_DELAY)
    self.char_delay = kwargs.get('char', DEFAULT_CHAR_DELAY)

  def _write_byte(self, val, mode):
    print("SENDING: " + bin(val)[2:].zfill(8) + ", MODE: " + str(mode))
    high_nib = (val >> 4) << 4
    low_nib  = (val & 0x0F) << 4
    print("  HIGH NIB: " + bin(high_nib | mode)[2:].zfill(8))
    self._i2c_write(high_nib | mode)
    print("  LOW NIB:  " + bin(low_nib | mode)[2:].zfill(8))
    self._i2c_write(low_nib | mode)
    self._pulse(self.backlight)

  def _i2c_write(self, val):
    data = val | self.backlight
    print("    SENDING BYTE: " + bin(data)[2:].zfill(8))
    self.bus.write_byte(self.i2c_addr, data)
    self._pulse(data)

  def _pulse(self, val):
    self.bus.write_byte(self.i2c_addr, val | 0x04)
    sleep(1/1000000)
    self.bus.write_byte(self.i2c_addr, val & 0xFB)
    sleep(50/1000000)

  # To avoid clashing with Python's print(), we'll break API compliance
  def printstr(self, val):
    # print the provided value on the display
    for c in val[:]:
      print("WRITTING: " + c)
      self.write(ord(c))
      sleep(5)

  def println(self, val):
    # Call print(), with trailing new line
    self.printstr(str(val) + "\n")

  def write(self, val):
    # raw write value to the display, callers are reponsible for implementing any delay after calling this method
    if not self.test:
      self._write_byte(val, LCD_REG_DATA)
      sleep(self.char_delay / 1000000)

  def command(self, val):
    # Send command to display, for display-specific commands not covered in the API
    self._write_byte(val, LCD_REG_CMD)
    sleep(self.cmd_delay / 1000000) # Pause to ensure command is executed
    return bin(val)[2:].zfill(8)

  def clear(self):
    # clear display and return cursor to 0,0
    return self.command(LCD_CMD_CLEARDISPLAY)

  def home(self):
    # set cursor position to 0,0 leaving display untouched
    return self.command(LCD_CMD_CURSORHOME)

  def set_cursor(self, row, col):
    # move cursor to indicated position, row and col values falling outside the configured
    # display size will be set to 0 or the rows/colums value provided in the constructor
    if row < 0:
      row = 0

    if col < 0:
      col = 0

    if row > self.rows:
      row = self.rows

    if col > self.cols:
      col = self.cols

    if row == 0 and col == 0:
      self.home()
    else:
      # TODO: move cursor to the indicated position.  Should we send cursor home command and then
      # move, or determine current cursor position and move relative to that?
      pass

  def cursor_on(self):
    # set block cursor on
    return self.command(LCD_CMD_DISPLAYCONTROL | LCD_CURSORON)

  def cursor_off(self):
    # set block cursor off
    return self.command(LCD_CMD_DISPLAYCONTROL | LCD_CURSOROFF)

  def blink_on(self):
    # set blinking underline cursor on
    return self.command(LCD_CMD_DISPLAYCONTROL | LCD_BLINKON)

  def blink_off(self):
    # set blinking underline cursor off
    return self.command(LCD_CMD_DISPLAYCONTROL | LCD_BLINKOFF)

  # Optional functions
  def set_backlight(self, val):
    # TODO: set backlight brightness (0-255), where 0 = off
    # Best I can tell, backlight control is either on or off
    if val > 0:
      self.backlight = LCD_BACKLIGHT
    else:
      self.backlight = LCD_NOBACKLIGHT

    self.bus.write_byte(self.i2c_addr, self.backlight)
    self.backlight

  def set_contrast(self, val):
    # TODO: set display contrast (0-255)
    pass

  def on(self):
    # turn display on and set backlight to ~ 75%
    self.command(LCD_CMD_DISPLAYCONTROL | LCD_DISPLAYON)
    return self.set_backlight(192)

  def off(self):
    # set backlight to 0 and turn display off
    self.set_backlight(0)
    return self.command(LCD_CMD_DISPLAYCONTROL | LCD_DISPLAYOFF)

  def status(self):
    # TODO: return status of the display. API docs say:
    #   Returns the FIFO buffer on the robot-electronics display
    #   Can be used to get r/w status of hardwired displays
    pass

  # Extended functions left unimplemented:
  #  load_custom_characters(self, char_num, rows)
  #  keypad(self)
  #  printstr(self)
  # These seem like a cool idea, and may implment in the future:
  #  init_bargraph(self, type)
  #  draw_horizontal_graph(row, col, len, end)
  #  draw_vertical_graph(row, col, len, end)

if __name__ == "__main__":
  # cls = hd44780_i2c(1, 0x3f, 20, 4, test = True)
  cls = hd44780_i2c(1, 0x3f, 20, 4)
  print("SLEEPING")
  sleep(3)
  cls.printstr('wad951')
#  print(cls.off())
#  print(cls.clear())
#  print(cls.blink_on())
<|endoftext|>