<commit_before>from datetime import datetime
from multiprocessing import Process
from time import time
from uuid import uuid4

from cbagent.collectors import NSServer
from cbagent.metadata_client import MetadataClient

from perfrunner.settings import CbAgentSettings

def with_stats(method):
    def wrapper(self, *args, **kwargs):
        self.cbagent.update_metadata()

        from_ts = self.cbagent.start()
        method(self, *args, **kwargs)
        to_ts = self.cbagent.stop()

        self.cbagent.add_snapshot(method.__name__, from_ts, to_ts)
    return wrapper

class CbAgent(object):

    def __init__(self, cluster_spec, target_iterator):
        settings = CbAgentSettings()
        settings.cluster = '{0}_{1}'.format(cluster_spec.name, uuid4().hex[:6])
        settings.ssh_username, settings.ssh_password = \
            cluster_spec.get_ssh_credentials()
        settings.rest_username, settings.rest_password = \
            cluster_spec.get_rest_credentials()

        self.md_client = MetadataClient(settings)
        self.snapshot_prefix = settings.cluster

        self.collectors = []
        for target in target_iterator:
            settings.master_node = target.node.split(':')[0]
            for collector in (NSServer, ):
                self.collectors.append(collector(settings))

    def update_metadata(self):
        for collector in self.collectors:
            collector.update_metadata()

    def start(self):
        self.processes = [Process(target=c.collect) for c in self.collectors]
        map(lambda p: p.start(), self.processes)
        return datetime.fromtimestamp(time())

    def stop(self):
        map(lambda p: p.terminate(), self.processes)
        return datetime.fromtimestamp(time())

    def add_snapshot(self, phase, ts_from, ts_to):
        name = '{0}_{1}'.format(self.snapshot_prefix, phase)
        self.md_client.add_snapshot(name, ts_from, ts_to)
<commit_msg>enable ActiveTasks stats collector<commit_after>from datetime import datetime
from multiprocessing import Process
from time import time
from uuid import uuid4

from cbagent.collectors import NSServer, ActiveTasks
from cbagent.metadata_client import MetadataClient

from perfrunner.settings import CbAgentSettings

def with_stats(method):
    def wrapper(self, *args, **kwargs):
        self.cbagent.update_metadata()

        from_ts = self.cbagent.start()
        method(self, *args, **kwargs)
        to_ts = self.cbagent.stop()

        self.cbagent.add_snapshot(method.__name__, from_ts, to_ts)
    return wrapper

class CbAgent(object):

    def __init__(self, cluster_spec, target_iterator):
        settings = CbAgentSettings()
        settings.cluster = '{0}_{1}'.format(cluster_spec.name, uuid4().hex[:6])
        settings.ssh_username, settings.ssh_password = \
            cluster_spec.get_ssh_credentials()
        settings.rest_username, settings.rest_password = \
            cluster_spec.get_rest_credentials()

        self.md_client = MetadataClient(settings)
        self.snapshot_prefix = settings.cluster

        self.collectors = []
        for target in target_iterator:
            settings.master_node = target.node.split(':')[0]
            for collector in (NSServer, ActiveTasks):
                self.collectors.append(collector(settings))

    def update_metadata(self):
        for collector in self.collectors:
            collector.update_metadata()

    def start(self):
        self.processes = [Process(target=c.collect) for c in self.collectors]
        map(lambda p: p.start(), self.processes)
        return datetime.fromtimestamp(time())

    def stop(self):
        map(lambda p: p.terminate(), self.processes)
        return datetime.fromtimestamp(time())

    def add_snapshot(self, phase, ts_from, ts_to):
        name = '{0}_{1}'.format(self.snapshot_prefix, phase)
        self.md_client.add_snapshot(name, ts_from, ts_to)
<|endoftext|>


<commit_before># Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
# For details: https://github.com/nedbat/coveragepy/blob/master/NOTICE.txt

"""Tests about understanding how third-party code is installed."""

import os
import os.path
import shutil

import pytest

from coverage import env

from tests.coveragetest import CoverageTest, TESTS_DIR
from tests.helpers import change_dir, make_file, nice_file
from tests.helpers import re_lines, run_command

def run_in_venv(cmd):
    r"""Run `cmd` in the virtualenv at `venv`.

    The first word of the command will be adjusted to run it from the
    venv/bin or venv\Scripts directory.

    Returns the text output of the command.
    """
    words = cmd.split()
    if env.WINDOWS:
        words[0] = fr"venv\Scripts\{words[0]}.exe"
    else:
        words[0] = fr"venv/bin/{words[0]}"
    status, output = run_command(" ".join(words))
    assert status == 0
    return output

@pytest.fixture(scope="session", name="venv_world")
def venv_world_fixture(tmp_path_factory):
    """Create a virtualenv with a few test packages for VirtualenvTest to use.

    Returns the directory containing the "venv" virtualenv.
    """

    venv_world = tmp_path_factory.mktemp("venv_world")
    with change_dir(venv_world):
        # Create a virtualenv.
        run_command("python -m venv venv")

        # A third-party package that installs a few different packages.
        make_file("third_pkg/third/__init__.py", """\
            import fourth
            def third(x):
                return 3 * x
            """)
        # Use plugin2.py as third.plugin
        with open(os.path.join(os.path.dirname(__file__), "plugin2.py")) as f:
            make_file("third_pkg/third/plugin.py", f.read())
        # A render function for plugin2 to use for dynamic file names.
        make_file("third_pkg/third/render.py", """\
            def render(filename, linenum):
                return "HTML: {}@{}".format(filename, linenum)
            """)
        # Another package that third can use.
        make_file("third_pkg/fourth/__init__.py", """\
            def fourth(x):
                return 4 * x
            """)
        # Some namespace packages.
        make_file("third_pkg/nspkg/fifth/__init__.py", """\
            def fifth(x):
                return 5 * x
            """)
        # The setup.py to install everything.
        make_file("third_pkg/setup.py", """\
            import setuptools
            setuptools.setup(
                name="third",
                packages=["third", "fourth", "nspkg.fifth"],
            )
            """)

        # Some namespace packages.
        make_file("another_pkg/nspkg/sixth/__init__.py", """\
            def sixth(x):
                return 6 * x
            """)
        make_file("another_pkg/setup.py", """\
            import setuptools
            setuptools.setup(
                name="another",
                packages=["nspkg.sixth"],
            )
            """)

        # Bug888 code.
        make_file("bug888/app/setup.py", """\
            from setuptools import setup
            setup(
                name='testcov',
                packages=['testcov'],
                namespace_packages=['testcov'],
            )
            """)
        make_file("bug888/app/testcov/__init__.py", """\
            try:  # pragma: no cover
                __import__('pkg_resources').declare_namespace(__name__)
            except ImportError:  # pragma: no cover
                from pkgutil import extend_path
                __path__ = extend_path(__path__, __name__)
            """)
        make_file("bug888/app/testcov/main.py", """\
            import pkg_resources
            for entry_point in pkg_resources.iter_entry_points('plugins'):
                entry_point.load()()
            """)
        make_file("bug888/plugin/setup.py", """\
            from setuptools import setup
            setup(
                name='testcov-plugin',
                packages=['testcov'],
                namespace_packages=['testcov'],
                entry_points={'plugins': ['testp = testcov.plugin:testp']},
            )
            """)
        make_file("bug888/plugin/testcov/__init__.py", """\
            try:  # pragma: no cover
                __import__('pkg_resources').declare_namespace(__name__)
            except ImportError:  # pragma: no cover
                from pkgutil import extend_path
                __path__ = extend_path(__path__, __name__)
            """)
        make_file("bug888/plugin/testcov/plugin.py", """\
            def testp():
                print("Plugin here")
            """)

        # Install the third-party packages.
        run_in_venv("python -m pip install --no-index ./third_pkg")
        run_in_venv("python -m pip install --no-index -e ./another_pkg")
        run_in_venv("python -m pip install --no-index -e ./bug888/app -e ./bug888/plugin")
        shutil.rmtree("third_pkg")

        # Install coverage.
        coverage_src = nice_file(TESTS_DIR, "..")
        run_in_venv(f"python -m pip install --no-index {coverage_src}")

    return venv_world

@pytest.fixture(params=[
    "coverage",
    "python -m coverage",
], name="coverage_command")
def coverage_command_fixture(request):
    """Parametrized fixture to use multiple forms of "coverage" command."""
    return request.param

class VirtualenvTest(CoverageTest):
    """Tests of virtualenv considerations."""

    expected_stdout = "33\n110\n198\n1.5\n"

    @pytest.fixture(autouse=True)
    def in_venv_world_fixture(self, venv_world):
        """For running tests inside venv_world, and cleaning up made files."""
        with change_dir(venv_world):
            self.make_file("myproduct.py", """\
                import colorsys
                import third
                import nspkg.fifth
                import nspkg.sixth
                print(third.third(11))
                print(nspkg.fifth.fifth(22))
                print(nspkg.sixth.sixth(33))
                print(sum(colorsys.rgb_to_hls(1, 0, 0)))
                """)

            self.del_environ("COVERAGE_TESTING")    # To avoid needing contracts installed.
            self.set_environ("COVERAGE_DEBUG_FILE", "debug_out.txt")
            self.set_environ("COVERAGE_DEBUG", "trace")

            yield

            for fname in os.listdir("."):
                if fname not in {"venv", "another_pkg", "bug888"}:
                    os.remove(fname)

    def get_trace_output(self):
        """Get the debug output of coverage.py"""
        with open("debug_out.txt") as f:
            return f.read()

    def test_third_party_venv_isnt_measured(self, coverage_command):
        out = run_in_venv(coverage_command + " run --source=. myproduct.py")
        # In particular, this warning doesn't appear:
        # Already imported a file that will be measured: .../coverage/__main__.py
        assert out == self.expected_stdout

        # Check that our tracing was accurate. Files are mentioned because
        # --source refers to a file.
        debug_out = self.get_trace_output()
        assert re_lines(
            r"^Not tracing .*\bexecfile.py': inside --source, but is third-party",
            debug_out,
            )
        assert re_lines(r"^Tracing .*\bmyproduct.py", debug_out)
        assert re_lines(
            r"^Not tracing .*\bcolorsys.py': falls outside the --source spec",
            debug_out,
            )

        out = run_in_venv(coverage_command + " report")
        assert "myproduct.py" in out
        assert "third" not in out
        assert "coverage" not in out
        assert "colorsys" not in out

    def test_us_in_venv_isnt_measured(self, coverage_command):
        out = run_in_venv(coverage_command + " run --source=third myproduct.py")
        assert out == self.expected_stdout

        # Check that our tracing was accurate. Modules are mentioned because
        # --source refers to a module.
        debug_out = self.get_trace_output()
        assert re_lines(
            r"^Not tracing .*\bexecfile.py': " +
            "module 'coverage.execfile' falls outside the --source spec",
            debug_out,
            )
        assert re_lines(
            r"^Not tracing .*\bmyproduct.py': module 'myproduct' falls outside the --source spec",
            debug_out,
            )
        assert re_lines(
            r"^Not tracing .*\bcolorsys.py': module 'colorsys' falls outside the --source spec",
            debug_out,
            )

        out = run_in_venv(coverage_command + " report")
        assert "myproduct.py" not in out
        assert "third" in out
        assert "coverage" not in out
        assert "colorsys" not in out

    def test_venv_isnt_measured(self, coverage_command):
        out = run_in_venv(coverage_command + " run myproduct.py")
        assert out == self.expected_stdout

        debug_out = self.get_trace_output()
        assert re_lines(r"^Not tracing .*\bexecfile.py': is part of coverage.py", debug_out)
        assert re_lines(r"^Tracing .*\bmyproduct.py", debug_out)
        assert re_lines(r"^Not tracing .*\bcolorsys.py': is in the stdlib", debug_out)

        out = run_in_venv(coverage_command + " report")
        assert "myproduct.py" in out
        assert "third" not in out
        assert "coverage" not in out
        assert "colorsys" not in out

    @pytest.mark.skipif(not env.C_TRACER, reason="Plugins are only supported with the C tracer.")
    def test_venv_with_dynamic_plugin(self, coverage_command):
        # https://github.com/nedbat/coveragepy/issues/1150
        # Django coverage plugin was incorrectly getting warnings:
        # "Already imported: ... django/template/blah.py"
        # It happened because coverage imported the plugin, which imported
        # Django, and then the Django files were reported as traceable.
        self.make_file(".coveragerc", "[run]\nplugins=third.plugin\n")
        self.make_file("myrender.py", """\
            import third.render
            print(third.render.render("hello.html", 1723))
            """)
        out = run_in_venv(coverage_command + " run --source=. myrender.py")
        # The output should not have this warning:
        # Already imported a file that will be measured: ...third/render.py (already-imported)
        assert out == "HTML: hello.html@1723\n"

    def test_installed_namespace_packages(self, coverage_command):
        # https://github.com/nedbat/coveragepy/issues/1231
        # When namespace packages were installed, they were considered
        # third-party packages.  Test that isn't still happening.
        out = run_in_venv(coverage_command + " run --source=nspkg myproduct.py")
        # In particular, this warning doesn't appear:
        # Already imported a file that will be measured: .../coverage/__main__.py
        assert out == self.expected_stdout

        # Check that our tracing was accurate. Files are mentioned because
        # --source refers to a file.
        debug_out = self.get_trace_output()
        assert re_lines(
            r"^Not tracing .*\bexecfile.py': " +
            "module 'coverage.execfile' falls outside the --source spec",
            debug_out,
            )
        assert re_lines(
            r"^Not tracing .*\bmyproduct.py': module 'myproduct' falls outside the --source spec",
            debug_out,
            )
        assert re_lines(
            r"^Not tracing .*\bcolorsys.py': module 'colorsys' falls outside the --source spec",
            debug_out,
            )

        out = run_in_venv(coverage_command + " report")

        # Name                                                       Stmts   Miss  Cover
        # ------------------------------------------------------------------------------
        # another_pkg/nspkg/sixth/__init__.py                            2      0   100%
        # venv/lib/python3.9/site-packages/nspkg/fifth/__init__.py       2      0   100%
        # ------------------------------------------------------------------------------
        # TOTAL                                                          4      0   100%

        assert "myproduct.py" not in out
        assert "third" not in out
        assert "coverage" not in out
        assert "colorsys" not in out
        assert "fifth" in out
        assert "sixth" in out

    def test_bug888(self, coverage_command):
        out = run_in_venv(
            coverage_command +
            " run --source=bug888/app,bug888/plugin bug888/app/testcov/main.py"
            )
        # When the test fails, the output includes "Already imported a file that will be measured"
        assert out == "Plugin here\n"
<commit_msg>perf(test): run pip just once, shaves 1.2s off the fixture<commit_after># Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
# For details: https://github.com/nedbat/coveragepy/blob/master/NOTICE.txt

"""Tests about understanding how third-party code is installed."""

import os
import os.path
import shutil

import pytest

from coverage import env

from tests.coveragetest import CoverageTest, TESTS_DIR
from tests.helpers import change_dir, make_file, nice_file
from tests.helpers import re_lines, run_command

def run_in_venv(cmd):
    r"""Run `cmd` in the virtualenv at `venv`.

    The first word of the command will be adjusted to run it from the
    venv/bin or venv\Scripts directory.

    Returns the text output of the command.
    """
    words = cmd.split()
    if env.WINDOWS:
        words[0] = fr"venv\Scripts\{words[0]}.exe"
    else:
        words[0] = fr"venv/bin/{words[0]}"
    status, output = run_command(" ".join(words))
    assert status == 0
    return output

@pytest.fixture(scope="session", name="venv_world")
def venv_world_fixture(tmp_path_factory):
    """Create a virtualenv with a few test packages for VirtualenvTest to use.

    Returns the directory containing the "venv" virtualenv.
    """

    venv_world = tmp_path_factory.mktemp("venv_world")
    with change_dir(venv_world):
        # Create a virtualenv.
        run_command("python -m venv venv")

        # A third-party package that installs a few different packages.
        make_file("third_pkg/third/__init__.py", """\
            import fourth
            def third(x):
                return 3 * x
            """)
        # Use plugin2.py as third.plugin
        with open(os.path.join(os.path.dirname(__file__), "plugin2.py")) as f:
            make_file("third_pkg/third/plugin.py", f.read())
        # A render function for plugin2 to use for dynamic file names.
        make_file("third_pkg/third/render.py", """\
            def render(filename, linenum):
                return "HTML: {}@{}".format(filename, linenum)
            """)
        # Another package that third can use.
        make_file("third_pkg/fourth/__init__.py", """\
            def fourth(x):
                return 4 * x
            """)
        # Some namespace packages.
        make_file("third_pkg/nspkg/fifth/__init__.py", """\
            def fifth(x):
                return 5 * x
            """)
        # The setup.py to install everything.
        make_file("third_pkg/setup.py", """\
            import setuptools
            setuptools.setup(
                name="third",
                packages=["third", "fourth", "nspkg.fifth"],
            )
            """)

        # Some namespace packages.
        make_file("another_pkg/nspkg/sixth/__init__.py", """\
            def sixth(x):
                return 6 * x
            """)
        make_file("another_pkg/setup.py", """\
            import setuptools
            setuptools.setup(
                name="another",
                packages=["nspkg.sixth"],
            )
            """)

        # Bug888 code.
        make_file("bug888/app/setup.py", """\
            from setuptools import setup
            setup(
                name='testcov',
                packages=['testcov'],
                namespace_packages=['testcov'],
            )
            """)
        make_file("bug888/app/testcov/__init__.py", """\
            try:  # pragma: no cover
                __import__('pkg_resources').declare_namespace(__name__)
            except ImportError:  # pragma: no cover
                from pkgutil import extend_path
                __path__ = extend_path(__path__, __name__)
            """)
        make_file("bug888/app/testcov/main.py", """\
            import pkg_resources
            for entry_point in pkg_resources.iter_entry_points('plugins'):
                entry_point.load()()
            """)
        make_file("bug888/plugin/setup.py", """\
            from setuptools import setup
            setup(
                name='testcov-plugin',
                packages=['testcov'],
                namespace_packages=['testcov'],
                entry_points={'plugins': ['testp = testcov.plugin:testp']},
            )
            """)
        make_file("bug888/plugin/testcov/__init__.py", """\
            try:  # pragma: no cover
                __import__('pkg_resources').declare_namespace(__name__)
            except ImportError:  # pragma: no cover
                from pkgutil import extend_path
                __path__ = extend_path(__path__, __name__)
            """)
        make_file("bug888/plugin/testcov/plugin.py", """\
            def testp():
                print("Plugin here")
            """)

        # Install everything.
        coverage_src = nice_file(TESTS_DIR, "..")
        run_in_venv(
            "python -m pip install --no-index " +
            "./third_pkg " +
            "-e ./another_pkg " +
            "-e ./bug888/app -e ./bug888/plugin " +
            coverage_src
        )
        shutil.rmtree("third_pkg")

    return venv_world

@pytest.fixture(params=[
    "coverage",
    "python -m coverage",
], name="coverage_command")
def coverage_command_fixture(request):
    """Parametrized fixture to use multiple forms of "coverage" command."""
    return request.param

class VirtualenvTest(CoverageTest):
    """Tests of virtualenv considerations."""

    expected_stdout = "33\n110\n198\n1.5\n"

    @pytest.fixture(autouse=True)
    def in_venv_world_fixture(self, venv_world):
        """For running tests inside venv_world, and cleaning up made files."""
        with change_dir(venv_world):
            self.make_file("myproduct.py", """\
                import colorsys
                import third
                import nspkg.fifth
                import nspkg.sixth
                print(third.third(11))
                print(nspkg.fifth.fifth(22))
                print(nspkg.sixth.sixth(33))
                print(sum(colorsys.rgb_to_hls(1, 0, 0)))
                """)

            self.del_environ("COVERAGE_TESTING")    # To avoid needing contracts installed.
            self.set_environ("COVERAGE_DEBUG_FILE", "debug_out.txt")
            self.set_environ("COVERAGE_DEBUG", "trace")

            yield

            for fname in os.listdir("."):
                if fname not in {"venv", "another_pkg", "bug888"}:
                    os.remove(fname)

    def get_trace_output(self):
        """Get the debug output of coverage.py"""
        with open("debug_out.txt") as f:
            return f.read()

    def test_third_party_venv_isnt_measured(self, coverage_command):
        out = run_in_venv(coverage_command + " run --source=. myproduct.py")
        # In particular, this warning doesn't appear:
        # Already imported a file that will be measured: .../coverage/__main__.py
        assert out == self.expected_stdout

        # Check that our tracing was accurate. Files are mentioned because
        # --source refers to a file.
        debug_out = self.get_trace_output()
        assert re_lines(
            r"^Not tracing .*\bexecfile.py': inside --source, but is third-party",
            debug_out,
            )
        assert re_lines(r"^Tracing .*\bmyproduct.py", debug_out)
        assert re_lines(
            r"^Not tracing .*\bcolorsys.py': falls outside the --source spec",
            debug_out,
            )

        out = run_in_venv(coverage_command + " report")
        assert "myproduct.py" in out
        assert "third" not in out
        assert "coverage" not in out
        assert "colorsys" not in out

    def test_us_in_venv_isnt_measured(self, coverage_command):
        out = run_in_venv(coverage_command + " run --source=third myproduct.py")
        assert out == self.expected_stdout

        # Check that our tracing was accurate. Modules are mentioned because
        # --source refers to a module.
        debug_out = self.get_trace_output()
        assert re_lines(
            r"^Not tracing .*\bexecfile.py': " +
            "module 'coverage.execfile' falls outside the --source spec",
            debug_out,
            )
        assert re_lines(
            r"^Not tracing .*\bmyproduct.py': module 'myproduct' falls outside the --source spec",
            debug_out,
            )
        assert re_lines(
            r"^Not tracing .*\bcolorsys.py': module 'colorsys' falls outside the --source spec",
            debug_out,
            )

        out = run_in_venv(coverage_command + " report")
        assert "myproduct.py" not in out
        assert "third" in out
        assert "coverage" not in out
        assert "colorsys" not in out

    def test_venv_isnt_measured(self, coverage_command):
        out = run_in_venv(coverage_command + " run myproduct.py")
        assert out == self.expected_stdout

        debug_out = self.get_trace_output()
        assert re_lines(r"^Not tracing .*\bexecfile.py': is part of coverage.py", debug_out)
        assert re_lines(r"^Tracing .*\bmyproduct.py", debug_out)
        assert re_lines(r"^Not tracing .*\bcolorsys.py': is in the stdlib", debug_out)

        out = run_in_venv(coverage_command + " report")
        assert "myproduct.py" in out
        assert "third" not in out
        assert "coverage" not in out
        assert "colorsys" not in out

    @pytest.mark.skipif(not env.C_TRACER, reason="Plugins are only supported with the C tracer.")
    def test_venv_with_dynamic_plugin(self, coverage_command):
        # https://github.com/nedbat/coveragepy/issues/1150
        # Django coverage plugin was incorrectly getting warnings:
        # "Already imported: ... django/template/blah.py"
        # It happened because coverage imported the plugin, which imported
        # Django, and then the Django files were reported as traceable.
        self.make_file(".coveragerc", "[run]\nplugins=third.plugin\n")
        self.make_file("myrender.py", """\
            import third.render
            print(third.render.render("hello.html", 1723))
            """)
        out = run_in_venv(coverage_command + " run --source=. myrender.py")
        # The output should not have this warning:
        # Already imported a file that will be measured: ...third/render.py (already-imported)
        assert out == "HTML: hello.html@1723\n"

    def test_installed_namespace_packages(self, coverage_command):
        # https://github.com/nedbat/coveragepy/issues/1231
        # When namespace packages were installed, they were considered
        # third-party packages.  Test that isn't still happening.
        out = run_in_venv(coverage_command + " run --source=nspkg myproduct.py")
        # In particular, this warning doesn't appear:
        # Already imported a file that will be measured: .../coverage/__main__.py
        assert out == self.expected_stdout

        # Check that our tracing was accurate. Files are mentioned because
        # --source refers to a file.
        debug_out = self.get_trace_output()
        assert re_lines(
            r"^Not tracing .*\bexecfile.py': " +
            "module 'coverage.execfile' falls outside the --source spec",
            debug_out,
            )
        assert re_lines(
            r"^Not tracing .*\bmyproduct.py': module 'myproduct' falls outside the --source spec",
            debug_out,
            )
        assert re_lines(
            r"^Not tracing .*\bcolorsys.py': module 'colorsys' falls outside the --source spec",
            debug_out,
            )

        out = run_in_venv(coverage_command + " report")

        # Name                                                       Stmts   Miss  Cover
        # ------------------------------------------------------------------------------
        # another_pkg/nspkg/sixth/__init__.py                            2      0   100%
        # venv/lib/python3.9/site-packages/nspkg/fifth/__init__.py       2      0   100%
        # ------------------------------------------------------------------------------
        # TOTAL                                                          4      0   100%

        assert "myproduct.py" not in out
        assert "third" not in out
        assert "coverage" not in out
        assert "colorsys" not in out
        assert "fifth" in out
        assert "sixth" in out

    def test_bug888(self, coverage_command):
        out = run_in_venv(
            coverage_command +
            " run --source=bug888/app,bug888/plugin bug888/app/testcov/main.py"
            )
        # When the test fails, the output includes "Already imported a file that will be measured"
        assert out == "Plugin here\n"
<|endoftext|>


<commit_before>import uuid
from datetime import timedelta
from decimal import Decimal

from django.conf import settings
from django.test import TestCase
from django.utils.six import binary_type

from .models import Test

class TestEncryptedField(TestCase):
    def test_settings_has_key(self):
        key = settings.CRYPTOGRAPHY_KEY
        self.assertIsNotNone(key)
        self.assertIsInstance(key, binary_type)

    def test_field_match(self):
        obj = Test(boolean=True, char='Hello, world!',
                   decimal=Decimal('1.00'), duration=timedelta(1),
                   float=42.0, integer=42, ip_addres='::1',
                   text='Too short...', uuid=uuid.uuid4())
        obj.encrypted_boolean = obj.boolean
        obj.encrypted_char = obj.char
        obj.encrypted_decimal = obj.decimal
        obj.encrypted_duration = obj.duration
        obj.encrypted_float = obj.float
        obj.encrypted_integer = obj.integer
        obj.encrypted_ip_addres = obj.ip_addres
        obj.encrypted_text = obj.text
        obj.encrypted_uuid = obj.uuid
        obj.save()
        del obj

        # Fetch the object
        obj = Test.objects.get()
        self.assertEqual(obj.boolean, obj.encrypted_boolean)
        self.assertEqual(obj.char, obj.encrypted_char)
        self.assertEqual(obj.date, obj.encrypted_date)
        self.assertEqual(obj.decimal, obj.encrypted_decimal)
        self.assertEqual(obj.duration, obj.encrypted_duration)
        self.assertEqual(obj.float, obj.encrypted_float)
        self.assertEqual(obj.integer, obj.encrypted_integer)
        self.assertEqual(obj.ip_addres, obj.encrypted_ip_addres)
        self.assertEqual(obj.text, obj.encrypted_text)
        self.assertEqual(obj.uuid, obj.encrypted_uuid)
<commit_msg>Added more unit tests for EncryptedField<commit_after>import uuid
from datetime import timedelta
from decimal import Decimal

from django.conf import settings
from django.db import models
from django.test import TestCase
from django.utils.six import binary_type
from django_cryptography.fields import EncryptedField

from .models import Test

class TestEncryptedField(TestCase):
    def test_settings_has_key(self):
        key = settings.CRYPTOGRAPHY_KEY
        self.assertIsNotNone(key)
        self.assertIsInstance(key, binary_type)

    def test_field_match(self):
        obj = Test(boolean=True, char='Hello, world!',
                   decimal=Decimal('1.00'), duration=timedelta(1),
                   float=42.0, integer=42, ip_addres='::1',
                   text='Too short...', uuid=uuid.uuid4())
        obj.encrypted_boolean = obj.boolean
        obj.encrypted_char = obj.char
        obj.encrypted_decimal = obj.decimal
        obj.encrypted_duration = obj.duration
        obj.encrypted_float = obj.float
        obj.encrypted_integer = obj.integer
        obj.encrypted_ip_addres = obj.ip_addres
        obj.encrypted_text = obj.text
        obj.encrypted_uuid = obj.uuid
        obj.save()
        del obj

        # Fetch the object
        obj = Test.objects.get()
        self.assertEqual(obj.boolean, obj.encrypted_boolean)
        self.assertEqual(obj.char, obj.encrypted_char)
        self.assertEqual(obj.date, obj.encrypted_date)
        self.assertEqual(obj.decimal, obj.encrypted_decimal)
        self.assertEqual(obj.duration, obj.encrypted_duration)
        self.assertEqual(obj.float, obj.encrypted_float)
        self.assertEqual(obj.integer, obj.encrypted_integer)
        self.assertEqual(obj.ip_addres, obj.encrypted_ip_addres)
        self.assertEqual(obj.text, obj.encrypted_text)
        self.assertEqual(obj.uuid, obj.encrypted_uuid)

    def test_field_checks(self):
        class BadField(models.Model):
            field = EncryptedField(models.CharField())

        model = BadField()
        errors = model.check()
        self.assertEqual(len(errors), 1)
        # The inner CharField is missing a max_length.
        self.assertEqual('encrypted.E001', errors[0].id)
        self.assertIn('max_length', errors[0].msg)

    def test_invalid_base_fields(self):
        class Related(models.Model):
            field = EncryptedField(models.ForeignKey('fields.Test'))

        obj = Related()
        errors = obj.check()

        self.assertEqual(1, len(errors))
        self.assertEqual('encrypted.E002', errors[0].id)
<|endoftext|>


<commit_before>import urllib2
import collections
import csv
import re
import datetime
import decimal

def convert_string(string):
    if '-' in string:
        return datetime.datetime.strptime(string, '%d-%b-%y')
    else:
        return decimal.Decimal(string)

_DIRECTIVES_PATTERN = re.compile('[\W_]+')
def historical_prices(symbol, start_date=None, end_date=None):
    params = ['q=%s' % symbol, 'output=csv']
    if start_date:
        params.append(start_date.strftime('startdate=%b+%d,+%Y'))
    if end_date:
        params.append(start_date.strftime('startdate=%b+%d,+%Y'))

    response = urllib2.urlopen('http://www.google.com/finance/historical?' + '&'.join(params))
    try:
        raw = csv.reader(response)

        tuple = collections.namedtuple('HistoricalPrice', [_DIRECTIVES_PATTERN.sub('', directive.lower()) for directive in next(raw)])
        return [tuple(*map(convert_string, row)) for row in raw]
    finally:
        response.close()
<commit_msg>Glaring typo in Google Finance API...<commit_after>import urllib2
import collections
import csv
import re
import datetime
import decimal

def convert_string(string):
    if '-' in string:
        return datetime.datetime.strptime(string, '%d-%b-%y')
    else:
        return decimal.Decimal(string)

_DIRECTIVES_PATTERN = re.compile('[\W_]+')
def historical_prices(symbol, start_date=None, end_date=None):
    params = ['q=%s' % symbol, 'output=csv']
    if start_date:
        params.append(start_date.strftime('startdate=%b+%d,+%Y'))
    if end_date:
        params.append(start_date.strftime('enddate=%b+%d,+%Y'))

    response = urllib2.urlopen('http://www.google.com/finance/historical?' + '&'.join(params))
    try:
        raw = csv.reader(response)

        tuple = collections.namedtuple('HistoricalPrice', [_DIRECTIVES_PATTERN.sub('', directive.lower()) for directive in next(raw)])
        return [tuple(*map(convert_string, row)) for row in raw]
    finally:
        response.close()
<|endoftext|>


<commit_before>"""CAS authentication backend"""

from urllib import urlencode, urlopen
from urlparse import urljoin

from django.conf import settings
from django.core.exceptions import ObjectDoesNotExist
from cas.models import User, Tgt, PgtIOU
from cas.utils import cas_response_callbacks

__all__ = ['CASBackend']

def _verify_cas1(ticket, service):
    """Verifies CAS 1.0 authentication ticket.

    Returns username on success and None on failure.
    """

    params = {'ticket': ticket, 'service': service}
    url = (urljoin(settings.CAS_SERVER_URL, 'validate') + '?' +
           urlencode(params))
    page = urlopen(url)
    try:
        verified = page.readline().strip()
        if verified == 'yes':
            return page.readline().strip()
        else:
            return None
    finally:
        page.close()

def _verify_cas2(ticket, service):
    """Verifies CAS 2.0+ XML-based authentication ticket.

    Returns username on success and None on failure.
    """

    try:
        from xml.etree import ElementTree
    except ImportError:
        from elementtree import ElementTree

    if settings.CAS_PROXY_CALLBACK:
        params = {'ticket': ticket, 'service': service, 'pgtUrl': settings.CAS_PROXY_CALLBACK}
    else:
        params = {'ticket': ticket, 'service': service}

    url = (urljoin(settings.CAS_SERVER_URL, 'proxyValidate') + '?' +
           urlencode(params))

    page = urlopen(url)
    try:
        response = page.read()
        tree = ElementTree.fromstring(response)

        #Useful for debugging
        #from xml.dom.minidom import parseString
        #from xml.etree import ElementTree
        #txt = ElementTree.tostring(tree)
        #print parseString(txt).toprettyxml()

        if tree[0].tag.endswith('authenticationSuccess'):
            if settings.CAS_RESPONSE_CALLBACKS:
                cas_response_callbacks(tree)
            return tree[0][0].text
        else:
            return None
    finally:
        page.close()

def verify_proxy_ticket(ticket, service):
    """Verifies CAS 2.0+ XML-based proxy ticket.

    Returns username on success and None on failure.
    """

    try:
        from xml.etree import ElementTree
    except ImportError:
        from elementtree import ElementTree

    params = {'ticket': ticket, 'service': service}

    url = (urljoin(settings.CAS_SERVER_URL, 'proxyValidate') + '?' +
           urlencode(params))

    page = urlopen(url)

    try:
        response = page.read()
        tree = ElementTree.fromstring(response)
        if tree[0].tag.endswith('authenticationSuccess'):
            username = tree[0][0].text
            proxies = []
            if len(tree[0]) > 1:
                for element in tree[0][1]:
                    proxies.append(element.text)
            return {"username": username, "proxies": proxies}
        else:
            return None
    finally:
        page.close()

_PROTOCOLS = {'1': _verify_cas1, '2': _verify_cas2}

if settings.CAS_VERSION not in _PROTOCOLS:
    raise ValueError('Unsupported CAS_VERSION %r' % settings.CAS_VERSION)

_verify = _PROTOCOLS[settings.CAS_VERSION]

class CASBackend(object):
    """CAS authentication backend"""

    supports_object_permissions = False
    supports_inactive_user = False

    def authenticate(self, ticket, service):
        """Verifies CAS ticket and gets or creates User object
            NB: Use of PT to identify proxy
        """

        username = _verify(ticket, service)
        if not username:
            return None
        try:
            user = User.objects.get(username__iexact=username)
        except User.DoesNotExist:
            # user will have an "unusable" password
            user = User.objects.create_user(username, '')
            user.save()
        return user

    def get_user(self, user_id):
        """Retrieve the user's entry in the User model if it exists"""

        try:
            return User.objects.get(pk=user_id)
        except User.DoesNotExist:
            return None
<commit_msg>Tickets: Adapting backends from the KTHse/django-cas2 fork.<commit_after>"""CAS authentication backend"""

import logging
from urllib import urlencode, urlopen
from urlparse import urljoin
from xml.dom import minidom

from django.conf import settings
from django.core.exceptions import ObjectDoesNotExist
from cas.exceptions import CasTicketException
from cas.models import User, Tgt, PgtIOU
from cas.utils import cas_response_callbacks

__all__ = ['CASBackend']

logger = logging.getLogger(__name__)

def _verify_cas1(ticket, service):
    """Verifies CAS 1.0 authentication ticket.

    Returns username on success and None on failure.
    """

    params = {'ticket': ticket, 'service': service}
    url = (urljoin(settings.CAS_SERVER_URL, 'validate') + '?' +
           urlencode(params))
    page = urlopen(url)
    try:
        verified = page.readline().strip()
        if verified == 'yes':
            return page.readline().strip()
        else:
            return None
    finally:
        page.close()

def _verify_cas2(ticket, service):
    """Verifies CAS 2.0+ XML-based authentication ticket.

    Returns username on success and None on failure.
    """

    try:
        from xml.etree import ElementTree
    except ImportError:
        from elementtree import ElementTree

    params = {'ticket': ticket, 'service': service}
    if settings.CAS_PROXY_CALLBACK:
        params['pgtUrl'] = settings.CAS_PROXY_CALLBACK

    url = (urljoin(settings.CAS_SERVER_URL, 'proxyValidate') + '?' +
           urlencode(params))

    page = urlopen(url)

    username = None

    try:
        response = page.read()
        tree = ElementTree.fromstring(response)
        document = minidom.parseString(response)

        #Useful for debugging
        #print document.toprettyxml()

        if tree[0].tag.endswith('authenticationSuccess'):
            if settings.CAS_RESPONSE_CALLBACKS:
                cas_response_callbacks(tree)

            username = tree[0][0].text

            pgt_el = document.getElementsByTagName('cas:proxyGrantingTicket')
            if pgt_el:
                pgt = pgt_el[0].firstChild.nodeValue
                try:
                    pgtIou = _get_pgtiou(pgt)
                    tgt = Tgt.objects.get(username=username)
                    tgt.tgt = pgtIou.tgt
                    tgt.save()
                    pgtIou.delete()
                except Tgt.DoesNotExist:
                    Tgt.objects.create(username=username, tgt=pgtIou.tgt)
                    pgtIou.delete()
                except Exception:
                    logger.error('Failed to do proxy authentication.')

        else:
            failure = document.getElementsByTagName('cas:authenticationFailure')
            if failure:
                logger.warn('Authentication failed from CAS server: %s',
                            failure[0].firstChild.nodeValue)

    except Exception as e:
        logger.error('Failed to verify CAS authentication', e)

    finally:
        page.close()

    return username

def verify_proxy_ticket(ticket, service):
    """Verifies CAS 2.0+ XML-based proxy ticket.

    Returns username on success and None on failure.
    """

    try:
        from xml.etree import ElementTree
    except ImportError:
        from elementtree import ElementTree

    params = {'ticket': ticket, 'service': service}

    url = (urljoin(settings.CAS_SERVER_URL, 'proxyValidate') + '?' +
           urlencode(params))

    page = urlopen(url)

    try:
        response = page.read()
        tree = ElementTree.fromstring(response)
        if tree[0].tag.endswith('authenticationSuccess'):
            username = tree[0][0].text
            proxies = []
            if len(tree[0]) > 1:
                for element in tree[0][1]:
                    proxies.append(element.text)
            return {"username": username, "proxies": proxies}
        else:
            return None
    finally:
        page.close()

_PROTOCOLS = {'1': _verify_cas1, '2': _verify_cas2}

if settings.CAS_VERSION not in _PROTOCOLS:
    raise ValueError('Unsupported CAS_VERSION %r' % settings.CAS_VERSION)

_verify = _PROTOCOLS[settings.CAS_VERSION]

def _get_pgtiou(pgt):
    """ Returns a PgtIOU object given a pgt.

        The PgtIOU (tgt) is set by the CAS server in a different request
        that has completed before this call, however, it may not be found in
        the database by this calling thread, hence the attempt to get the
        ticket is retried for up to 5 seconds. This should be handled some
        better way.
    """
    pgtIou = None
    retries_left = 5
    while not pgtIou and retries_left:
        try:
            return PgtIOU.objects.get(tgt=pgt)
        except PgtIOU.DoesNotExist:
            time.sleep(1)
            retries_left -= 1
    raise CasTicketException("Could not find pgtIou for pgt %s" % pgt)

class CASBackend(object):
    """CAS authentication backend"""

    supports_object_permissions = False
    supports_inactive_user = False

    def authenticate(self, ticket, service):
        """Verifies CAS ticket and gets or creates User object
            NB: Use of PT to identify proxy
        """

        username = _verify(ticket, service)
        if not username:
            return None
        try:
            user = User.objects.get(username__iexact=username)
        except User.DoesNotExist:
            # user will have an "unusable" password
            user = User.objects.create_user(username, '')
            user.save()
        return user

    def get_user(self, user_id):
        """Retrieve the user's entry in the User model if it exists"""

        try:
            return User.objects.get(pk=user_id)
        except User.DoesNotExist:
            return None
<|endoftext|>


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

from luma.oled.device import ssd1306, sh1106
from luma.core.render import canvas

from PIL import ImageDraw, ImageFont, Image

import urllib2
import json
import time

from daemon import runner

class StratuxScreen():
    def __init__(self):
        self.stdin_path = '/dev/null'
        self.stdout_path = '/var/log/stratux-screen.log'
        self.stderr_path = '/var/log/stratux-screen.log'
        self.pidfile_path = '/var/run/stratux-screen.pid'
        self.pidfile_timeout = 5
    def run(self):
        font2 = ImageFont.truetype('/etc/stratux-screen/CnC_Red_Alert.ttf', 12)
        oled = ssd1306(port=1, address=0x3C)

        with canvas(oled) as draw:
            logo = Image.open('/etc/stratux-screen/stratux-logo-64x64.bmp')
            draw.bitmap((32, 0), logo, fill=1)

        time.sleep(10)
        n = 0

        while 1:
            time.sleep(1)
            response = urllib2.urlopen('http://localhost/getStatus')
            getStatusHTML = response.read()
            getStatusData = json.loads(getStatusHTML)
            CPUTemp = getStatusData["CPUTemp"]
            uat_current = getStatusData["UAT_messages_last_minute"]
            uat_max = getStatusData["UAT_messages_max"]
            es_current = getStatusData["ES_messages_last_minute"]
            es_max = getStatusData["ES_messages_max"]

            response = urllib2.urlopen('http://localhost/getTowers')
            getTowersHTML = response.read()
            getTowersData = json.loads(getTowersHTML)
            NumTowers = 0
            for towerLatLng in getTowersData:
                print getTowersData[towerLatLng]["Messages_last_minute"]
                if (getTowersData[towerLatLng]["Messages_last_minute"] > 0):
                    NumTowers += 1

            with canvas(oled) as draw:
                pad = 2 # Two pixels on the left and right.
                text_margin = 25
                # UAT status.
                draw.text((50, 0), "UAT", font=font2, fill=255)
                # "Status bar", 2 pixels high.
                status_bar_width_max = oled.width - (2 * pad) - (2 * text_margin)
                status_bar_width = 0
                if uat_max > 0:
                    status_bar_width = int((float(uat_current) / uat_max) * status_bar_width_max)
                draw.rectangle((pad + text_margin, 14, pad + text_margin + status_bar_width, 20), outline=255, fill=255) # Top left, bottom right.
                # Draw the current (left) and max (right) numbers.
                draw.text((pad, 14), str(uat_current), font=font2, fill=255)
                draw.text(((2*pad) + text_margin + status_bar_width_max, 14), str(uat_max), font=font2, fill=255)
                # ES status.
                draw.text((44, 24), "1090ES", font=font2, fill=255)
                status_bar_width = 0
                if es_max > 0:
                    status_bar_width = int((float(es_current) / es_max) * status_bar_width_max)
                draw.rectangle((pad + text_margin, 34, pad + text_margin + status_bar_width, 40), outline=255, fill=255) # Top left, bottom right.
                # Draw the current (left) and max (right) numbers.
                draw.text((pad, 34), str(es_current), font=font2, fill=255)
                draw.text(((2*pad) + text_margin + status_bar_width_max, 34), str(es_max), font=font2, fill=255)
                # Other stats.
                seq = (n / 5) % 2
                t = ""
                if seq == 0:
                    t = "CPU: %0.1fC, Towers: %d" % (CPUTemp, NumTowers)
                if seq == 1:
                    t = "GPS Sat: %d/%d/%d" % (getStatusData["GPS_satellites_locked"], getStatusData["GPS_satellites_seen"], getStatusData["GPS_satellites_tracked"])
                    if getStatusData["GPS_solution"] == "GPS + SBAS (WAAS / EGNOS)":
                        t = t + " (WAAS)"
                #print t
                draw.text((pad, 45), t, font=font2, fill=255)

            n = n+1

stratuxscreen = StratuxScreen()
daemon_runner = runner.DaemonRunner(stratuxscreen)
daemon_runner.do_action()
<commit_msg>Change luma usage.<commit_after>#!/usr/bin/env python

from luma.core.interface.serial import i2c
from luma.oled.device import ssd1306, sh1106
from luma.core.render import canvas

from PIL import ImageDraw, ImageFont, Image

import urllib2
import json
import time

from daemon import runner

class StratuxScreen():
    def __init__(self):
        self.stdin_path = '/dev/null'
        self.stdout_path = '/var/log/stratux-screen.log'
        self.stderr_path = '/var/log/stratux-screen.log'
        self.pidfile_path = '/var/run/stratux-screen.pid'
        self.pidfile_timeout = 5
    def run(self):
        font2 = ImageFont.truetype('/etc/stratux-screen/CnC_Red_Alert.ttf', 12)
        serial = i2c(port=1, address=0x3c)
        oled = ssd1306(serial)

        with canvas(oled) as draw:
            logo = Image.open('/etc/stratux-screen/stratux-logo-64x64.bmp')
            draw.bitmap((32, 0), logo, fill=1)

        time.sleep(10)
        n = 0

        while 1:
            time.sleep(1)
            response = urllib2.urlopen('http://localhost/getStatus')
            getStatusHTML = response.read()
            getStatusData = json.loads(getStatusHTML)
            CPUTemp = getStatusData["CPUTemp"]
            uat_current = getStatusData["UAT_messages_last_minute"]
            uat_max = getStatusData["UAT_messages_max"]
            es_current = getStatusData["ES_messages_last_minute"]
            es_max = getStatusData["ES_messages_max"]

            response = urllib2.urlopen('http://localhost/getTowers')
            getTowersHTML = response.read()
            getTowersData = json.loads(getTowersHTML)
            NumTowers = 0
            for towerLatLng in getTowersData:
                print getTowersData[towerLatLng]["Messages_last_minute"]
                if (getTowersData[towerLatLng]["Messages_last_minute"] > 0):
                    NumTowers += 1

            with canvas(oled) as draw:
                pad = 2 # Two pixels on the left and right.
                text_margin = 25
                # UAT status.
                draw.text((50, 0), "UAT", font=font2, fill=255)
                # "Status bar", 2 pixels high.
                status_bar_width_max = oled.width - (2 * pad) - (2 * text_margin)
                status_bar_width = 0
                if uat_max > 0:
                    status_bar_width = int((float(uat_current) / uat_max) * status_bar_width_max)
                draw.rectangle((pad + text_margin, 14, pad + text_margin + status_bar_width, 20), outline=255, fill=255) # Top left, bottom right.
                # Draw the current (left) and max (right) numbers.
                draw.text((pad, 14), str(uat_current), font=font2, fill=255)
                draw.text(((2*pad) + text_margin + status_bar_width_max, 14), str(uat_max), font=font2, fill=255)
                # ES status.
                draw.text((44, 24), "1090ES", font=font2, fill=255)
                status_bar_width = 0
                if es_max > 0:
                    status_bar_width = int((float(es_current) / es_max) * status_bar_width_max)
                draw.rectangle((pad + text_margin, 34, pad + text_margin + status_bar_width, 40), outline=255, fill=255) # Top left, bottom right.
                # Draw the current (left) and max (right) numbers.
                draw.text((pad, 34), str(es_current), font=font2, fill=255)
                draw.text(((2*pad) + text_margin + status_bar_width_max, 34), str(es_max), font=font2, fill=255)
                # Other stats.
                seq = (n / 5) % 2
                t = ""
                if seq == 0:
                    t = "CPU: %0.1fC, Towers: %d" % (CPUTemp, NumTowers)
                if seq == 1:
                    t = "GPS Sat: %d/%d/%d" % (getStatusData["GPS_satellites_locked"], getStatusData["GPS_satellites_seen"], getStatusData["GPS_satellites_tracked"])
                    if getStatusData["GPS_solution"] == "GPS + SBAS (WAAS / EGNOS)":
                        t = t + " (WAAS)"
                #print t
                draw.text((pad, 45), t, font=font2, fill=255)

            n = n+1

stratuxscreen = StratuxScreen()
daemon_runner = runner.DaemonRunner(stratuxscreen)
daemon_runner.do_action()
<|endoftext|>


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

import unittest
from werkzeug.exceptions import NotFound, Forbidden
from flask import Response
from jinja2 import TemplateNotFound
from coaster.views import jsonp
from nodular import Node, NodeView, NodePublisher, NodeRegistry
from .test_db import db, TestDatabaseFixture
from .test_nodetree import TestType

class MyNodeView(NodeView):
    @NodeView.route('/')
    def index(self):
        return u"node-index"

    def dummy(self):
        pass

class ExpandedNodeView(MyNodeView):
    @NodeView.route('/', methods=['GET', 'POST'])
    def index(self):
        # This never gets called for GET because MyNodeView.index is registered first
        return u'expanded-index'

    @NodeView.route('/edit', methods=['GET'])
    def editget(self):
        return u'edit-GET'

    @NodeView.route('/edit', methods=['POST'])
    def editpost(self):
        return u'edit-POST'

    @NodeView.route('/multimethod', methods=['GET', 'POST'])
    @NodeView.route('/multimethod', methods=['PUT'])
    def multimethod(self):
        from flask import request
        return u'multimethod-' + request.method

class RestrictedView(MyNodeView):
    @NodeView.route('/view')
    @NodeView.requires_permission('view')
    def view(self):
        return u'view'

    @NodeView.route('/admin')
    @NodeView.requires_permission('admin')
    def admin(self):
        return u'admin'

def viewcallable(data):
    return Response(repr(data), mimetype='text/plain')

class RenderedView(NodeView):
    @NodeView.route('/renderedview1')
    @NodeView.render_with('renderedview1.html')
    def myview(self):
        return {'data': 'value'}

    @NodeView.route('/renderedview2')
    @NodeView.render_with({
        'text/html': 'renderedview2.html',
        'text/xml': 'renderedview2.xml',
        'text/plain': viewcallable})
    def otherview(self):
        return {'data': 'value'}

# --- Tests -------------------------------------------------------------------

class TestNodeView(unittest.TestCase):
    def test_view_route(self):
        """Test that NodeView.route defines a route."""
        self.assertTrue(hasattr(MyNodeView, 'url_map'))
        self.assertEqual(type(MyNodeView.index), type(MyNodeView.dummy))
        self.assertEqual(len(list(MyNodeView.url_map.iter_rules())), 1)

class TestNodeViews(TestDatabaseFixture):
    def setUp(self):
        super(TestNodeViews, self).setUp()

        self.registry = NodeRegistry()
        self.registry.register_node(Node, view=MyNodeView, child_nodetypes=['*'])
        self.registry.register_node(TestType, view=MyNodeView, child_nodetypes=['*'], parent_nodetypes=['*'])

        self.registry.register_view('node', ExpandedNodeView)
        self.registry.register_view('test_type', ExpandedNodeView)

        # Make some nodes
        self.root = Node(name=u'root', title=u'Root Node')
        if not hasattr(self, 'nodetype'):
            self.nodetype = Node
        self.node1 = self.nodetype(name=u'node1', title=u'Node 1', parent=self.root)
        self.node2 = self.nodetype(name=u'node2', title=u'Node 2', parent=self.root)
        self.node3 = self.nodetype(name=u'node3', title=u'Node 3', parent=self.node2)
        self.node4 = self.nodetype(name=u'node4', title=u'Node 4', parent=self.node3)
        self.node5 = self.nodetype(name=u'node5', title=u'Node 5', parent=self.root)
        db.session.add_all([self.root, self.node1, self.node2, self.node3, self.node4, self.node5])
        db.session.commit()

        self.rootpub = NodePublisher(self.registry, u'/')
        self.nodepub = NodePublisher(self.registry, u'/node2', u'/')

    def test_publishview(self):
        """Publish a default view."""
        with self.app.test_request_context():
            response = self.rootpub.publish(u'/node2')
        self.assertEqual(response, 'node-index')

    def test_methods(self):
        """Publish views with different methods."""
        with self.app.test_request_context(method='GET'):
            response = self.rootpub.publish(u'/node2/edit')
        self.assertEqual(response, 'edit-GET')

        with self.app.test_request_context(method='POST'):
            response = self.rootpub.publish(u'/node2/edit')
        self.assertEqual(response, 'edit-POST')

        with self.app.test_request_context(method='GET'):
            response = self.rootpub.publish(u'/node2/multimethod')
        self.assertEqual(response, 'multimethod-GET')

        with self.app.test_request_context(method='POST'):
            response = self.rootpub.publish(u'/node2/multimethod')
        self.assertEqual(response, 'multimethod-POST')

        with self.app.test_request_context(method='PUT'):
            response = self.rootpub.publish(u'/node2/multimethod')
        self.assertEqual(response, 'multimethod-PUT')

        with self.app.test_request_context(method='GET'):
            self.assertRaises(NotFound, self.rootpub.publish, u'/node2/random')

class TestTypeViews(TestNodeViews):
    def setUp(self):
        self.nodetype = TestType
        super(TestTypeViews, self).setUp()

class TestPermissionAndRenderedViews(TestDatabaseFixture):
    def setUp(self):
        super(TestPermissionAndRenderedViews, self).setUp()

        self.registry = NodeRegistry()
        self.registry.register_node(TestType, view=RestrictedView)
        self.registry.register_view('test_type', RenderedView)
        self.root = Node(name=u'root', title=u'Root Node')
        self.node = TestType(name=u'node', title=u'Node', parent=self.root)
        db.session.add_all([self.root, self.node])
        db.session.commit()
        self.publisher = NodePublisher(self.registry, u'/')

    def test_view(self):
        """
        Test access to the restricted view.
        """
        # No permission required to access '/'
        with self.app.test_request_context(method='GET'):
            response = self.publisher.publish(u'/node')
        self.assertEqual(response, u'node-index')
        # 'view' permission is granted to everyone on TestType
        with self.app.test_request_context(method='GET'):
            response = self.publisher.publish(u'/node/view')
        self.assertEqual(response, u'view')
        # 'admin' permission is granted to no one on TestType
        with self.app.test_request_context(method='GET'):
            self.assertRaises(Forbidden, self.publisher.publish, u'/node/admin',
                user=self.user1, permissions=['siteadmin'])
        # Calling without a request context gets you the actual value
        view = RenderedView(self.node)
        result = view.myview()
        self.assertEqual(result, {'data': 'value'})

    def test_render(self):
        """
        Test rendered views.
        """
        # For this test to pass, the render_view decorator must call render_template
        # with the correct template name. Since the templates don't actually exist,
        # we'll get a TemplateNotFound exception, so our "test" is to confirm that the
        # missing template is the one that was supposed to be rendered.
        with self.app.test_request_context(method='GET', headers=[('Accept', '')]):
            try:
                self.publisher.publish(u'/node/renderedview1')
            except TemplateNotFound, e:
                self.assertEqual(str(e), 'renderedview1.html')
            else:
                raise Exception("Wrong template rendered")

        for acceptheader, template in [
                ('text/html,text/xml,*/*', 'renderedview2.html'),
                ('text/xml,text/html,*/*', 'renderedview2.xml')]:
            with self.app.test_request_context(method='GET', headers=[('Accept', acceptheader)]):
                try:
                    self.publisher.publish(u'/node/renderedview2')
                except TemplateNotFound, e:
                    self.assertEqual(str(e), template)
                else:
                    raise Exception("Wrong template rendered")

        # The application/json and text/plain renderers do exist, so we should get
        # a valid return value from them.
        with self.app.test_request_context(method='GET', headers=[('Accept', 'application/json')]):
            response = self.publisher.publish(u'/node/renderedview2')
            self.assertTrue(isinstance(response, Response))
            self.assertEqual(response.data, jsonp({"data": "value"}).data)
        with self.app.test_request_context(method='GET', headers=[('Accept', 'text/plain')]):
            response = self.publisher.publish(u'/node/renderedview2')
            self.assertTrue(isinstance(response, Response))
            self.assertEqual(response.data, "{'data': 'value'}")
<commit_msg>Improve test coverage.<commit_after># -*- coding: utf-8 -*-

import unittest
from werkzeug.exceptions import NotFound, Forbidden
from flask import Response
from jinja2 import TemplateNotFound
from coaster.views import jsonp
from nodular import Node, NodeView, NodePublisher, NodeRegistry
from .test_db import db, TestDatabaseFixture
from .test_nodetree import TestType

class MyNodeView(NodeView):
    @NodeView.route('/')
    def index(self):
        return u"node-index"

    def dummy(self):
        pass

class ExpandedNodeView(MyNodeView):
    @NodeView.route('/', methods=['GET', 'POST'])
    def index(self):
        # This never gets called for GET because MyNodeView.index is registered first
        return u'expanded-index'

    @NodeView.route('/edit', methods=['GET'])
    def editget(self):
        return u'edit-GET'

    @NodeView.route('/edit', methods=['POST'])
    def editpost(self):
        return u'edit-POST'

    @NodeView.route('/multimethod', methods=['GET', 'POST'])
    @NodeView.route('/multimethod', methods=['PUT'])
    def multimethod(self):
        from flask import request
        return u'multimethod-' + request.method

class RestrictedView(MyNodeView):
    @NodeView.route('/view')
    @NodeView.requires_permission('view')
    def view(self):
        return u'view'

    @NodeView.route('/admin')
    @NodeView.requires_permission('admin')
    def admin(self):
        return u'admin'

def viewcallable(data):
    return Response(repr(data), mimetype='text/plain')

class RenderedView(NodeView):
    @NodeView.route('/renderedview1')
    @NodeView.render_with('renderedview1.html')
    def myview(self):
        return {'data': 'value'}

    @NodeView.route('/renderedview2')
    @NodeView.render_with({
        'text/html': 'renderedview2.html',
        'text/xml': 'renderedview2.xml',
        'text/plain': viewcallable})
    def otherview(self):
        return {'data': 'value'}

# --- Tests -------------------------------------------------------------------

class TestNodeView(unittest.TestCase):
    def test_view_route(self):
        """Test that NodeView.route defines a route."""
        self.assertTrue(hasattr(MyNodeView, 'url_map'))
        self.assertEqual(type(MyNodeView.index), type(MyNodeView.dummy))
        self.assertEqual(len(list(MyNodeView.url_map.iter_rules())), 1)

class TestNodeViews(TestDatabaseFixture):
    def setUp(self):
        super(TestNodeViews, self).setUp()

        self.registry = NodeRegistry()
        self.registry.register_node(Node, view=MyNodeView, child_nodetypes=['*'])
        self.registry.register_node(TestType, view=MyNodeView, child_nodetypes=['*'], parent_nodetypes=['*'])

        self.registry.register_view('node', ExpandedNodeView)
        self.registry.register_view('test_type', ExpandedNodeView)

        # Make some nodes
        self.root = Node(name=u'root', title=u'Root Node')
        if not hasattr(self, 'nodetype'):
            self.nodetype = Node
        self.node1 = self.nodetype(name=u'node1', title=u'Node 1', parent=self.root)
        self.node2 = self.nodetype(name=u'node2', title=u'Node 2', parent=self.root)
        self.node3 = self.nodetype(name=u'node3', title=u'Node 3', parent=self.node2)
        self.node4 = self.nodetype(name=u'node4', title=u'Node 4', parent=self.node3)
        self.node5 = self.nodetype(name=u'node5', title=u'Node 5', parent=self.root)
        db.session.add_all([self.root, self.node1, self.node2, self.node3, self.node4, self.node5])
        db.session.commit()

        self.rootpub = NodePublisher(self.registry, u'/')
        self.nodepub = NodePublisher(self.registry, u'/node2', u'/')

    def test_publishview(self):
        """Publish a default view."""
        with self.app.test_request_context():
            response = self.rootpub.publish(u'/node2')
        self.assertEqual(response, 'node-index')

    def test_methods(self):
        """Publish views with different methods."""
        with self.app.test_request_context(method='GET'):
            response = self.rootpub.publish(u'/node2/edit')
        self.assertEqual(response, 'edit-GET')

        with self.app.test_request_context(method='POST'):
            response = self.rootpub.publish(u'/node2/edit')
        self.assertEqual(response, 'edit-POST')

        with self.app.test_request_context(method='GET'):
            response = self.rootpub.publish(u'/node2/multimethod')
        self.assertEqual(response, 'multimethod-GET')

        with self.app.test_request_context(method='POST'):
            response = self.rootpub.publish(u'/node2/multimethod')
        self.assertEqual(response, 'multimethod-POST')

        with self.app.test_request_context(method='PUT'):
            response = self.rootpub.publish(u'/node2/multimethod')
        self.assertEqual(response, 'multimethod-PUT')

        with self.app.test_request_context(method='GET'):
            self.assertRaises(NotFound, self.rootpub.publish, u'/node2/random')

class TestTypeViews(TestNodeViews):
    def setUp(self):
        self.nodetype = TestType
        super(TestTypeViews, self).setUp()

class TestPermissionAndRenderedViews(TestDatabaseFixture):
    def setUp(self):
        super(TestPermissionAndRenderedViews, self).setUp()

        self.registry = NodeRegistry()
        self.registry.register_node(TestType, view=RestrictedView)
        self.registry.register_view(TestType, RenderedView)
        self.root = Node(name=u'root', title=u'Root Node')
        self.node = TestType(name=u'node', title=u'Node', parent=self.root)
        db.session.add_all([self.root, self.node])
        db.session.commit()
        self.publisher = NodePublisher(self.registry, u'/')

    def test_view(self):
        """
        Test access to the restricted view.
        """
        # No permission required to access '/'
        with self.app.test_request_context(method='GET'):
            response = self.publisher.publish(u'/node')
        self.assertEqual(response, u'node-index')
        # 'view' permission is granted to everyone on TestType
        with self.app.test_request_context(method='GET'):
            response = self.publisher.publish(u'/node/view')
        self.assertEqual(response, u'view')
        # 'admin' permission is granted to no one on TestType
        with self.app.test_request_context(method='GET'):
            self.assertRaises(Forbidden, self.publisher.publish, u'/node/admin',
                user=self.user1, permissions=['siteadmin'])
        # Calling without a request context gets you the actual value
        view = RenderedView(self.node)
        result = view.myview()
        self.assertEqual(result, {'data': 'value'})

    def test_render(self):
        """
        Test rendered views.
        """
        # For this test to pass, the render_view decorator must call render_template
        # with the correct template name. Since the templates don't actually exist,
        # we'll get a TemplateNotFound exception, so our "test" is to confirm that the
        # missing template is the one that was supposed to be rendered.
        with self.app.test_request_context(method='GET', headers=[('Accept', '')]):
            try:
                self.publisher.publish(u'/node/renderedview1')
            except TemplateNotFound, e:
                self.assertEqual(str(e), 'renderedview1.html')
            else:
                raise Exception("Wrong template rendered")

        for acceptheader, template in [
                ('text/html,text/xml,*/*', 'renderedview2.html'),
                ('text/xml,text/html,*/*', 'renderedview2.xml')]:
            with self.app.test_request_context(method='GET', headers=[('Accept', acceptheader)]):
                try:
                    self.publisher.publish(u'/node/renderedview2')
                except TemplateNotFound, e:
                    self.assertEqual(str(e), template)
                else:
                    raise Exception("Wrong template rendered")

        # The application/json and text/plain renderers do exist, so we should get
        # a valid return value from them.
        with self.app.test_request_context(method='GET', headers=[('Accept', 'application/json')]):
            response = self.publisher.publish(u'/node/renderedview2')
            self.assertTrue(isinstance(response, Response))
            self.assertEqual(response.data, jsonp({"data": "value"}).data)
        with self.app.test_request_context(method='GET', headers=[('Accept', 'text/plain')]):
            response = self.publisher.publish(u'/node/renderedview2')
            self.assertTrue(isinstance(response, Response))
            self.assertEqual(response.data, "{'data': 'value'}")
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright 2018 Data61, CSIRO
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Graph node classification using GraphSAGE.
This currently is only tested on the CORA dataset, which can be downloaded from https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz

The following is the description of the dataset:
> The Cora dataset consists of 2708 scientific publications classified into one of seven classes.
> The citation network consists of 5429 links. Each publication in the dataset is described by a
> 0/1-valued word vector indicating the absence/presence of the corresponding word from the dictionary.
> The dictionary consists of 1433 unique words. The README file in the dataset provides more details.

Download and unzip the cora.tgz file to a location on your computer and pass this location
(which should contain cora.cites and cora.content) as a command line argument to this script.

Run this script as follows:
    python graphsage-cora-example.py -l <path_to_cora_dataset>

Other optional arguments can be seen by running
    python graphsage-cora-example.py --help

"""
import os
import argparse
import pickle
import numpy as np
import pandas as pd
from scipy import sparse
import networkx as nx
import keras
from keras import optimizers, losses, layers, metrics
from keras.callbacks import EarlyStopping, TensorBoard, ModelCheckpoint
from sklearn import preprocessing, feature_extraction, model_selection
import stellargraph as sg
from stellargraph.layer import GAT
from stellargraph.mapper import FullBatchNodeGenerator

def train(
    edgelist,
    node_data,
    attn_heads,
    layer_sizes,
    num_epochs=10,
    learning_rate=0.005,
    es_patience=100,
    dropout=0.0,
    target_name="subject",
):
    """
    Train a GAT model on the specified graph G with given parameters, evaluate it, and save the model.

    Args:
        edgelist: Graph edgelist
        node_data: Feature and target data for nodes
        attn_heads: Number of attention heads in GAT layers
        layer_sizes: A list of number of hidden nodes in each layer
        num_epochs: Number of epochs to train the model
        learning_rate: Initial Learning rate
        dropout: The dropout (0->1)
    """
    # Extract target and encode as a one-hot vector
    target_encoding = feature_extraction.DictVectorizer(sparse=False)
    node_targets = target_encoding.fit_transform(
        node_data[[target_name]].to_dict("records")
    )
    node_ids = node_data.index

    # Extract the feature data. These are the feature vectors that the Keras model will use as input.
    # The CORA dataset contains attributes 'w_x' that correspond to words found in that publication.
    node_features = node_data[feature_names]

    # Create graph from edgelist and set node features and node type
    Gnx = nx.from_pandas_edgelist(edgelist)

    # Convert to StellarGraph and prepare for ML
    G = sg.StellarGraph(Gnx, node_type_name="label", node_features=node_features)

    # Split nodes into train/test using stratification.
    train_nodes, test_nodes, train_targets, test_targets = model_selection.train_test_split(
        node_ids,
        node_targets,
        train_size=140,
        test_size=None,
        stratify=node_targets,
        random_state=55232,
    )

    # Further split test set into validation and test
    val_nodes, test_nodes, val_targets, test_targets = model_selection.train_test_split(
        test_nodes, test_targets, train_size=500, test_size=1000, random_state=523214
    )

    # Create mappers for GraphSAGE that input data from the graph to the model
    generator = FullBatchNodeGenerator(G)
    train_gen = generator.flow(train_nodes, train_targets)
    val_gen = generator.flow(val_nodes, val_targets)

    # GAT model
    gat = GAT(
        layer_sizes=layer_sizes,
        attn_heads=attn_heads,
        generator=train_gen,
        bias=True,
        in_dropout=dropout,
        attn_dropout=dropout,
        activations=["elu", "elu"],
        normalize=None,
    )
    # Expose the input and output tensors of the GAT model:
    x_inp, x_out = gat.node_model()

    # Snap the final estimator layer to x_out
    x_out = layers.Dense(units=train_targets.shape[1], activation="softmax")(x_out)

    # Create Keras model for training
    model = keras.Model(inputs=x_inp, outputs=x_out)
    model.compile(
        optimizer=optimizers.Adam(lr=learning_rate, decay=0.001),
        loss=losses.categorical_crossentropy,
        weighted_metrics=["acc"],
    )
    print(model.summary())

    # Train model
    # Callbacks
    if not os.path.isdir("logs"):
        os.makedirs("logs")
    es_callback = EarlyStopping(monitor="val_weighted_acc", patience=es_patience)
    mc_callback = ModelCheckpoint(
        "logs/best_model.h5",
        monitor="val_weighted_acc",
        save_best_only=True,
        save_weights_only=True,
    )

    if args.interface == "fit":
        print("\nUsing model.fit() to train the model\n")
        # Get the training data
        [X, A], y_train, node_mask_train = train_gen.__getitem__(0)
        N = A.shape[0]
        # A = sparse.csr_matrix(A + np.eye(A.shape[0]))  # Add self-loops

        # Get the validation data
        [_, _], y_val, node_mask_val = val_gen.__getitem__(0)

        history = model.fit(
            x=[X, A],
            y=y_train,
            sample_weight=node_mask_train,
            batch_size=N,
            shuffle=False,  # must be False, since shuffling data means shuffling the whole graph
            epochs=num_epochs,
            verbose=2,
            validation_data=([X, A], y_val, node_mask_val),
            callbacks=[es_callback, mc_callback],
        )
    else:
        print("\nUsing model.fit_generator() to train the model\n")
        history = model.fit_generator(
            train_gen,
            epochs=num_epochs,
            validation_data=val_gen,
            verbose=2,
            shuffle=False,
            callbacks=[es_callback, mc_callback],
        )

    # Load best model
    model.load_weights("logs/best_model.h5")

    # Evaluate on validation set and print metrics
    if args.interface == "fit":
        val_metrics = model.evaluate(
            x=[X, A], y=y_val, sample_weight=node_mask_val, batch_size=N
        )
    else:
        val_metrics = model.evaluate_generator(val_gen)

    print("\nBest model's Validation Set Metrics:")
    for name, val in zip(model.metrics_names, val_metrics):
        print("\t{}: {:0.4f}".format(name, val))

    # Evaluate on test set and print metrics
    if args.interface == "fit":
        [_, _], y_test, node_mask_test = generator.flow(
            test_nodes, test_targets
        ).__getitem__(0)
        test_metrics = model.evaluate(
            x=[X, A], y=y_test, sample_weight=node_mask_test, batch_size=N
        )
    else:
        test_metrics = model.evaluate_generator(
            generator.flow(test_nodes, test_targets)
        )

    print("\nBest model's Test Set Metrics:")
    for name, val in zip(model.metrics_names, test_metrics):
        print("\t{}: {:0.4f}".format(name, val))

    # Get predictions for all nodes
    if args.interface == "fit":
        all_predictions = model.predict(x=[X, A], batch_size=N)
    else:
        all_predictions = model.predict_generator(generator.flow(node_ids))

    # Turn predictions back into the original categories
    node_predictions = pd.DataFrame(
        target_encoding.inverse_transform(all_predictions), index=list(G.nodes())
    )
    accuracy = np.mean(
        [
            "subject=" + gt_subject == p
            for gt_subject, p in zip(
                node_data["subject"][list(G.nodes())], node_predictions.idxmax(axis=1)
            )
        ]
    )
    print("\nAll-node accuracy: {:0.4f}".format(accuracy))

    # TODO: extract the GAT embeddings from x_out, and save/plot them

    # Save the trained model
    save_str = "_h{}_l{}_d{}_r{}".format(
        attn_heads, "_".join([str(x) for x in layer_sizes]), dropout, learning_rate
    )
    model.save("cora_gat_model" + save_str + ".h5")

    # We must also save the target encoding to convert model predictions
    with open("cora_gat_encoding" + save_str + ".pkl", "wb") as f:
        pickle.dump([target_encoding], f)

def test():
    raise NotImplemented

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Graph node classification using Graph Attention Network (GAT)"
    )
    parser.add_argument(
        "-c",
        "--checkpoint",
        nargs="?",
        type=str,
        default=None,
        help="Load a saved checkpoint .h5 file",
    )
    parser.add_argument(
        "-e",
        "--epochs",
        type=int,
        default=10,
        help="The number of epochs to train the model",
    )
    parser.add_argument(
        "-d",
        "--dropout",
        type=float,
        default=0.5,
        help="Dropout rate for the GAT model, between 0.0 and 1.0",
    )
    parser.add_argument(
        "-r",
        "--learningrate",
        type=float,
        default=0.005,
        help="Initial learning rate for model training",
    )
    parser.add_argument(
        "-p", "--patience", type=int, default=10,
        help="Patience for early stopping (number of epochs with no improvement after which training should be stopped)"
    )
    parser.add_argument(
        "-a", "--attn_heads", type=int, default=1, help="Number of attention heads"
    )
    parser.add_argument(
        "-s",
        "--layer_sizes",
        type=int,
        nargs="*",
        default=[8, 8],
        help="The number of hidden features at each GAT layer",
    )
    parser.add_argument(
        "-l",
        "--location",
        type=str,
        default=None,
        help="Location of the CORA dataset (directory)",
    )
    parser.add_argument(
        "-t",
        "--target",
        type=str,
        default="subject",
        help="The target node attribute (categorical)",
    )
    parser.add_argument(
        "-i",
        "--interface",
        type=str,
        default="fit_generator",
        help="Defines which method is used for model training (.fit() or .fit_generator())",
    )
    args, cmdline_args = parser.parse_known_args()

    # Load the dataset - this assumes it is the CORA dataset
    # Load graph edgelist
    if args.location is not None:
        graph_loc = os.path.expanduser(args.location)
    else:
        raise ValueError(
            "Please specify the directory containing the dataset using the '-l' flag"
        )

    edgelist = pd.read_table(
        os.path.join(graph_loc, "cora.cites"), header=None, names=["source", "target"]
    )

    # Load node features
    # The CORA dataset contains binary attributes 'w_x' that correspond to whether the corresponding keyword
    # (out of 1433 keywords) is found in the corresponding publication.
    feature_names = ["w_{}".format(ii) for ii in range(1433)]
    # Also, there is a "subject" column
    column_names = feature_names + ["subject"]
    node_data = pd.read_table(
        os.path.join(graph_loc, "cora.content"), header=None, names=column_names
    )

    if args.checkpoint is None:
        train(
            edgelist,
            node_data,
            args.attn_heads,
            args.layer_sizes,
            args.epochs,
            args.learningrate,
            args.patience,
            args.dropout,
        )
    else:
        test()
<commit_msg>added TensorBoard callback<commit_after># -*- coding: utf-8 -*-
#
# Copyright 2018 Data61, CSIRO
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Graph node classification using GraphSAGE.
This currently is only tested on the CORA dataset, which can be downloaded from https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz

The following is the description of the dataset:
> The Cora dataset consists of 2708 scientific publications classified into one of seven classes.
> The citation network consists of 5429 links. Each publication in the dataset is described by a
> 0/1-valued word vector indicating the absence/presence of the corresponding word from the dictionary.
> The dictionary consists of 1433 unique words. The README file in the dataset provides more details.

Download and unzip the cora.tgz file to a location on your computer and pass this location
(which should contain cora.cites and cora.content) as a command line argument to this script.

Run this script as follows:
    python graphsage-cora-example.py -l <path_to_cora_dataset>

Other optional arguments can be seen by running
    python graphsage-cora-example.py --help

"""
import os
import argparse
import pickle
import numpy as np
import pandas as pd
from scipy import sparse
import networkx as nx
import keras
from keras import optimizers, losses, layers, metrics
from keras.callbacks import EarlyStopping, TensorBoard, ModelCheckpoint
from sklearn import preprocessing, feature_extraction, model_selection
import stellargraph as sg
from stellargraph.layer import GAT
from stellargraph.mapper import FullBatchNodeGenerator

def train(
    edgelist,
    node_data,
    attn_heads,
    layer_sizes,
    num_epochs=10,
    learning_rate=0.005,
    es_patience=100,
    dropout=0.0,
    target_name="subject",
):
    """
    Train a GAT model on the specified graph G with given parameters, evaluate it, and save the model.

    Args:
        edgelist: Graph edgelist
        node_data: Feature and target data for nodes
        attn_heads: Number of attention heads in GAT layers
        layer_sizes: A list of number of hidden nodes in each layer
        num_epochs: Number of epochs to train the model
        learning_rate: Initial Learning rate
        dropout: The dropout (0->1)
    """
    # Extract target and encode as a one-hot vector
    target_encoding = feature_extraction.DictVectorizer(sparse=False)
    node_targets = target_encoding.fit_transform(
        node_data[[target_name]].to_dict("records")
    )
    node_ids = node_data.index

    # Extract the feature data. These are the feature vectors that the Keras model will use as input.
    # The CORA dataset contains attributes 'w_x' that correspond to words found in that publication.
    node_features = node_data[feature_names]

    # Create graph from edgelist and set node features and node type
    Gnx = nx.from_pandas_edgelist(edgelist)

    # Convert to StellarGraph and prepare for ML
    G = sg.StellarGraph(Gnx, node_type_name="label", node_features=node_features)

    # Split nodes into train/test using stratification.
    train_nodes, test_nodes, train_targets, test_targets = model_selection.train_test_split(
        node_ids,
        node_targets,
        train_size=140,
        test_size=None,
        stratify=node_targets,
        random_state=55232,
    )

    # Further split test set into validation and test
    val_nodes, test_nodes, val_targets, test_targets = model_selection.train_test_split(
        test_nodes, test_targets, train_size=500, test_size=1000, random_state=523214
    )

    # Create mappers for GraphSAGE that input data from the graph to the model
    generator = FullBatchNodeGenerator(G)
    train_gen = generator.flow(train_nodes, train_targets)
    val_gen = generator.flow(val_nodes, val_targets)

    # GAT model
    gat = GAT(
        layer_sizes=layer_sizes,
        attn_heads=attn_heads,
        generator=train_gen,
        bias=True,
        in_dropout=dropout,
        attn_dropout=dropout,
        activations=["elu", "elu"],
        normalize=None,
    )
    # Expose the input and output tensors of the GAT model:
    x_inp, x_out = gat.node_model()

    # Snap the final estimator layer to x_out
    x_out = layers.Dense(units=train_targets.shape[1], activation="softmax")(x_out)

    # Create Keras model for training
    model = keras.Model(inputs=x_inp, outputs=x_out)
    model.compile(
        optimizer=optimizers.Adam(lr=learning_rate, decay=0.001),
        loss=losses.categorical_crossentropy,
        weighted_metrics=["acc"],
    )
    print(model.summary())

    # Train model
    # Callbacks
    if not os.path.isdir("logs"):
        os.makedirs("logs")
    N = len(node_ids)
    es_callback = EarlyStopping(monitor="val_weighted_acc", patience=es_patience)
    tb_callback = TensorBoard(batch_size=N)
    mc_callback = ModelCheckpoint(
        "logs/best_model.h5",
        monitor="val_weighted_acc",
        save_best_only=True,
        save_weights_only=True,
    )

    if args.interface == "fit":
        print("\nUsing model.fit() to train the model\n")
        # Get the training data
        [X, A], y_train, node_mask_train = train_gen.__getitem__(0)
        N = A.shape[0]
        # A = sparse.csr_matrix(A + np.eye(A.shape[0]))  # Add self-loops

        # Get the validation data
        [_, _], y_val, node_mask_val = val_gen.__getitem__(0)

        history = model.fit(
            x=[X, A],
            y=y_train,
            sample_weight=node_mask_train,
            batch_size=N,
            shuffle=False,  # must be False, since shuffling data means shuffling the whole graph
            epochs=num_epochs,
            verbose=2,
            validation_data=([X, A], y_val, node_mask_val),
            callbacks=[es_callback, tb_callback, mc_callback],
        )
    else:
        print("\nUsing model.fit_generator() to train the model\n")
        history = model.fit_generator(
            train_gen,
            epochs=num_epochs,
            validation_data=val_gen,
            verbose=2,
            shuffle=False,
            callbacks=[es_callback, tb_callback, mc_callback],
        )

    # Load best model
    model.load_weights("logs/best_model.h5")

    # Evaluate on validation set and print metrics
    if args.interface == "fit":
        val_metrics = model.evaluate(
            x=[X, A], y=y_val, sample_weight=node_mask_val, batch_size=N
        )
    else:
        val_metrics = model.evaluate_generator(val_gen)

    print("\nBest model's Validation Set Metrics:")
    for name, val in zip(model.metrics_names, val_metrics):
        print("\t{}: {:0.4f}".format(name, val))

    # Evaluate on test set and print metrics
    if args.interface == "fit":
        [_, _], y_test, node_mask_test = generator.flow(
            test_nodes, test_targets
        ).__getitem__(0)
        test_metrics = model.evaluate(
            x=[X, A], y=y_test, sample_weight=node_mask_test, batch_size=N
        )
    else:
        test_metrics = model.evaluate_generator(
            generator.flow(test_nodes, test_targets)
        )

    print("\nBest model's Test Set Metrics:")
    for name, val in zip(model.metrics_names, test_metrics):
        print("\t{}: {:0.4f}".format(name, val))

    # Get predictions for all nodes
    if args.interface == "fit":
        all_predictions = model.predict(x=[X, A], batch_size=N)
    else:
        all_predictions = model.predict_generator(generator.flow(node_ids))

    # Turn predictions back into the original categories
    node_predictions = pd.DataFrame(
        target_encoding.inverse_transform(all_predictions), index=list(G.nodes())
    )
    accuracy = np.mean(
        [
            "subject=" + gt_subject == p
            for gt_subject, p in zip(
                node_data["subject"][list(G.nodes())], node_predictions.idxmax(axis=1)
            )
        ]
    )
    print("\nAll-node accuracy: {:0.4f}".format(accuracy))

    # TODO: extract the GAT embeddings from x_out, and save/plot them

    # Save the trained model
    save_str = "_h{}_l{}_d{}_r{}".format(
        attn_heads, "_".join([str(x) for x in layer_sizes]), dropout, learning_rate
    )
    model.save("cora_gat_model" + save_str + ".h5")

    # We must also save the target encoding to convert model predictions
    with open("cora_gat_encoding" + save_str + ".pkl", "wb") as f:
        pickle.dump([target_encoding], f)

def test():
    raise NotImplemented

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Graph node classification using Graph Attention Network (GAT)"
    )
    parser.add_argument(
        "-c",
        "--checkpoint",
        nargs="?",
        type=str,
        default=None,
        help="Load a saved checkpoint .h5 file",
    )
    parser.add_argument(
        "-e",
        "--epochs",
        type=int,
        default=10,
        help="The number of epochs to train the model",
    )
    parser.add_argument(
        "-d",
        "--dropout",
        type=float,
        default=0.5,
        help="Dropout rate for the GAT model, between 0.0 and 1.0",
    )
    parser.add_argument(
        "-r",
        "--learningrate",
        type=float,
        default=0.005,
        help="Initial learning rate for model training",
    )
    parser.add_argument(
        "-p", "--patience", type=int, default=10,
        help="Patience for early stopping (number of epochs with no improvement after which training should be stopped)"
    )
    parser.add_argument(
        "-a", "--attn_heads", type=int, default=1, help="Number of attention heads"
    )
    parser.add_argument(
        "-s",
        "--layer_sizes",
        type=int,
        nargs="*",
        default=[8, 8],
        help="The number of hidden features at each GAT layer",
    )
    parser.add_argument(
        "-l",
        "--location",
        type=str,
        default=None,
        help="Location of the CORA dataset (directory)",
    )
    parser.add_argument(
        "-t",
        "--target",
        type=str,
        default="subject",
        help="The target node attribute (categorical)",
    )
    parser.add_argument(
        "-i",
        "--interface",
        type=str,
        default="fit_generator",
        help="Defines which method is used for model training (.fit() or .fit_generator())",
    )
    args, cmdline_args = parser.parse_known_args()

    # Load the dataset - this assumes it is the CORA dataset
    # Load graph edgelist
    if args.location is not None:
        graph_loc = os.path.expanduser(args.location)
    else:
        raise ValueError(
            "Please specify the directory containing the dataset using the '-l' flag"
        )

    edgelist = pd.read_table(
        os.path.join(graph_loc, "cora.cites"), header=None, names=["source", "target"]
    )

    # Load node features
    # The CORA dataset contains binary attributes 'w_x' that correspond to whether the corresponding keyword
    # (out of 1433 keywords) is found in the corresponding publication.
    feature_names = ["w_{}".format(ii) for ii in range(1433)]
    # Also, there is a "subject" column
    column_names = feature_names + ["subject"]
    node_data = pd.read_table(
        os.path.join(graph_loc, "cora.content"), header=None, names=column_names
    )

    if args.checkpoint is None:
        train(
            edgelist,
            node_data,
            args.attn_heads,
            args.layer_sizes,
            args.epochs,
            args.learningrate,
            args.patience,
            args.dropout,
        )
    else:
        test()
<|endoftext|>