<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import logging
import re
import sys
from os.path import exists as path_exists

import pytest

from pyscaffold.api import create_project
from pyscaffold.cli import parse_args, run
from pyscaffold.extensions import django
from pyscaffold.templates import setup_py

skip_py33 = pytest.mark.skipif(sys.version_info[:2] == (3, 3),
                               reason="django-admin.py fails with Python 3.3")

PROJ_NAME = "proj"
DJANGO_FILES = ["proj/manage.py", "proj/proj/wsgi.py"]

@skip_py33
@pytest.mark.slow
def test_create_project_with_django(tmpfolder):
    # Given options with the django extension,
    opts = dict(project=PROJ_NAME, extensions=[django.Django('django')])

    # when the project is created,
    create_project(opts)

    # then django files should exist
    for path in DJANGO_FILES:
        assert path_exists(path)
    # and also overwritable pyscaffold files (with the exact contents)
    tmpfolder.join(PROJ_NAME).join("setup.py").read() == setup_py(opts)

@skip_py33
def test_pretend_create_project_with_django(tmpfolder, caplog):
    # Given options with the django extension,
    caplog.set_level(logging.INFO)
    opts = parse_args([PROJ_NAME, '--pretend', '--django'])

    # when the project is created,
    create_project(opts)

    # then files should exist
    assert not path_exists(PROJ_NAME)
    for path in DJANGO_FILES:
        assert not path_exists(path)

    # but activities should be logged
    logs = caplog.text
    assert re.search(r'run.+django', logs)

def test_create_project_without_django(tmpfolder):
    # Given options without the django extension,
    opts = dict(project=PROJ_NAME)

    # when the project is created,
    create_project(opts)

    # then django files should not exist
    for path in DJANGO_FILES:
        assert not path_exists(path)

def test_create_project_no_django(tmpfolder, nodjango_admin_mock):
    # Given options with the django extension,
    # but without django-admin being installed,
    opts = dict(project=PROJ_NAME, extensions=[django.Django('django')])

    # when the project is created,
    # then an exception should be raised.
    with pytest.raises(django.DjangoAdminNotInstalled):
        create_project(opts)

@skip_py33
@pytest.mark.slow
def test_cli_with_django(tmpfolder):
    # Given the command line with the django option,
    sys.argv = ["pyscaffold", "--django", PROJ_NAME]

    # when pyscaffold runs,
    run()

    # then django files should exist
    for path in DJANGO_FILES:
        assert path_exists(path)

def test_cli_without_django(tmpfolder):
    # Given the command line without the django option,
    sys.argv = ["pyscaffold", PROJ_NAME]

    # when pyscaffold runs,
    run()

    # then django files should not exist
    for path in DJANGO_FILES:
        assert not path_exists(path)

def test_cli_with_django_and_update(tmpfolder, capsys):
    # Given a project exists
    create_project(project=PROJ_NAME)

    # when the project is updated
    # with the django extension,
    sys.argv = ["pyscaffold", PROJ_NAME, "--update", "--django"]
    run()

    # then a warning should be displayed
    out, err = capsys.readouterr()
    assert all(warn in out + err for warn in (
        'external tools', 'not supported', 'will be ignored'))
<commit_msg>Change unittest<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import logging
import re
import sys
from os.path import exists as path_exists

import pytest

from pyscaffold.api import create_project
from pyscaffold.cli import parse_args, run
from pyscaffold.extensions import django
from pyscaffold.templates import setup_py

skip_py33 = pytest.mark.skipif(sys.version_info[:2] == (3, 3),
                               reason="django-admin.py fails with Python 3.3")

PROJ_NAME = "proj"
DJANGO_FILES = ["proj/manage.py", "proj/src/proj/wsgi.py"]

@skip_py33
@pytest.mark.slow
def test_create_project_with_django(tmpfolder):
    # Given options with the django extension,
    opts = dict(project=PROJ_NAME, extensions=[django.Django('django')])

    # when the project is created,
    create_project(opts)

    # then django files should exist
    for path in DJANGO_FILES:
        assert path_exists(path)
    # and also overwritable pyscaffold files (with the exact contents)
    tmpfolder.join(PROJ_NAME).join("setup.py").read() == setup_py(opts)

@skip_py33
def test_pretend_create_project_with_django(tmpfolder, caplog):
    # Given options with the django extension,
    caplog.set_level(logging.INFO)
    opts = parse_args([PROJ_NAME, '--pretend', '--django'])

    # when the project is created,
    create_project(opts)

    # then files should exist
    assert not path_exists(PROJ_NAME)
    for path in DJANGO_FILES:
        assert not path_exists(path)

    # but activities should be logged
    logs = caplog.text
    assert re.search(r'run.+django', logs)

def test_create_project_without_django(tmpfolder):
    # Given options without the django extension,
    opts = dict(project=PROJ_NAME)

    # when the project is created,
    create_project(opts)

    # then django files should not exist
    for path in DJANGO_FILES:
        assert not path_exists(path)

def test_create_project_no_django(tmpfolder, nodjango_admin_mock):
    # Given options with the django extension,
    # but without django-admin being installed,
    opts = dict(project=PROJ_NAME, extensions=[django.Django('django')])

    # when the project is created,
    # then an exception should be raised.
    with pytest.raises(django.DjangoAdminNotInstalled):
        create_project(opts)

@skip_py33
@pytest.mark.slow
def test_cli_with_django(tmpfolder):
    # Given the command line with the django option,
    sys.argv = ["pyscaffold", "--django", PROJ_NAME]

    # when pyscaffold runs,
    run()

    # then django files should exist
    for path in DJANGO_FILES:
        assert path_exists(path)

def test_cli_without_django(tmpfolder):
    # Given the command line without the django option,
    sys.argv = ["pyscaffold", PROJ_NAME]

    # when pyscaffold runs,
    run()

    # then django files should not exist
    for path in DJANGO_FILES:
        assert not path_exists(path)

def test_cli_with_django_and_update(tmpfolder, capsys):
    # Given a project exists
    create_project(project=PROJ_NAME)

    # when the project is updated
    # with the django extension,
    sys.argv = ["pyscaffold", PROJ_NAME, "--update", "--django"]
    run()

    # then a warning should be displayed
    out, err = capsys.readouterr()
    assert all(warn in out + err for warn in (
        'external tools', 'not supported', 'will be ignored'))
<|endoftext|>


<commit_before>import sys
import re
import textwrap
from doctest import OutputChecker, ELLIPSIS
from tests.lib import reset_env, run_pip, write_file, get_env, pyversion, pip_install_local
from tests.lib.local_repos import local_checkout, local_repo

distribute_re = re.compile('^distribute==[0-9.]+\n', re.MULTILINE)

def _check_output(result, expected):
    checker = OutputChecker()
    actual = str(result)

    ## FIXME!  The following is a TOTAL hack.  For some reason the
    ## __str__ result for pkg_resources.Requirement gets downcased on
    ## Windows.  Since INITools is the only package we're installing
    ## in this file with funky case requirements, I'm forcibly
    ## upcasing it.  You can also normalize everything to lowercase,
    ## but then you have to remember to upcase <BLANKLINE>.  The right
    ## thing to do in the end is probably to find out how to report
    ## the proper fully-cased package name in our error message.
    if sys.platform == 'win32':
        actual = actual.replace('initools', 'INITools')

    # This allows our existing tests to work when run in a context
    # with distribute installed.
    actual = distribute_re.sub('', actual)

    def banner(msg):
        return '\n========== %s ==========\n' % msg
    assert checker.check_output(expected, actual, ELLIPSIS), banner('EXPECTED')+expected+banner('ACTUAL')+actual+banner(6*'=')

def test_freeze_basic():
    """
    Some tests of freeze, first we have to install some stuff.  Note that
    the test is a little crude at the end because Python 2.5+ adds egg
    info to the standard library, so stuff like wsgiref will show up in
    the freezing.  (Probably that should be accounted for in pip, but
    currently it is not).

    """
    env = reset_env()
    write_file('initools-req.txt', textwrap.dedent("""\
        simple==2.0
        # and something else to test out:
        simple2<=3.0
        """))
    result = pip_install_local('-r', env.scratch_path/'initools-req.txt')
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: pip freeze
        -- stdout: --------------------
        ...simple==2.0
        simple2==3.0...
        <BLANKLINE>""")
    _check_output(result, expected)

def test_freeze_svn():
    """Test freezing a svn checkout"""

    checkout_path = local_checkout('svn+http://svn.colorstudy.com/INITools/trunk')
    #svn internally stores windows drives as uppercase; we'll match that.
    checkout_path = checkout_path.replace('c:', 'C:')

    env = reset_env()
    result = env.run('svn', 'co', '-r10',
                     local_repo('svn+http://svn.colorstudy.com/INITools/trunk'),
                     'initools-trunk')
    result = env.run('python', 'setup.py', 'develop',
            cwd=env.scratch_path/ 'initools-trunk', expect_stderr=True)
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        -e %s@10#egg=INITools-0.3.1dev...-dev_r10
        ...""" % checkout_path)
    _check_output(result, expected)

def test_freeze_git_clone():
    """
    Test freezing a Git clone.

    """
    env = reset_env()
    result = env.run('git', 'clone', local_repo('git+http://github.com/pypa/pip-test-package.git'), 'pip-test-package')
    result = env.run('git', 'checkout', '7d654e66c8fa7149c165ddeffa5b56bc06619458',
            cwd=env.scratch_path / 'pip-test-package', expect_stderr=True)
    result = env.run('python', 'setup.py', 'develop',
            cwd=env.scratch_path / 'pip-test-package')
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        ...-e %s@...#egg=pip_test_package-...
        ...""" % local_checkout('git+http://github.com/pypa/pip-test-package.git'))
    _check_output(result, expected)

    result = run_pip('freeze', '-f',
                     '%s#egg=pip_test_package' % local_checkout('git+http://github.com/pypa/pip-test-package.git'),
                     expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: pip freeze -f %(repo)s#egg=pip_test_package
        -- stdout: --------------------
        -f %(repo)s#egg=pip_test_package...
        -e %(repo)s@...#egg=pip_test_package-0.1.1
        ...""" % {'repo': local_checkout('git+http://github.com/pypa/pip-test-package.git')})
    _check_output(result, expected)

def test_freeze_mercurial_clone():
    """
    Test freezing a Mercurial clone.

    """
    reset_env()
    env = get_env()
    result = env.run('hg', 'clone',
                     '-r', '7bc186caa7dc',
                     local_repo('hg+http://bitbucket.org/jezdez/django-authority'),
                     'django-authority')
    result = env.run('python', 'setup.py', 'develop',
            cwd=env.scratch_path/'django-authority', expect_stderr=True)
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        ...-e %s@...#egg=django_authority-...
        ...""" % local_checkout('hg+http://bitbucket.org/jezdez/django-authority'))
    _check_output(result, expected)

    result = run_pip('freeze', '-f',
                     '%s#egg=django_authority' % local_checkout('hg+http://bitbucket.org/jezdez/django-authority'),
                     expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze -f %(repo)s#egg=django_authority
        -- stdout: --------------------
        -f %(repo)s#egg=django_authority
        ...-e %(repo)s@...#egg=django_authority-dev
        ...""" % {'repo': local_checkout('hg+http://bitbucket.org/jezdez/django-authority')})
    _check_output(result, expected)

def test_freeze_bazaar_clone():
    """
    Test freezing a Bazaar clone.

    """

    checkout_path = local_checkout('bzr+http://bazaar.launchpad.net/%7Edjango-wikiapp/django-wikiapp/release-0.1')
    #bzr internally stores windows drives as uppercase; we'll match that.
    checkout_pathC = checkout_path.replace('c:', 'C:')

    reset_env()
    env = get_env()
    result = env.run('bzr', 'checkout', '-r', '174',
                     local_repo('bzr+http://bazaar.launchpad.net/%7Edjango-wikiapp/django-wikiapp/release-0.1'),
                     'django-wikiapp')
    result = env.run('python', 'setup.py', 'develop',
            cwd=env.scratch_path/'django-wikiapp')
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        ...-e %s@...#egg=django_wikiapp-...
        ...""" % checkout_pathC)
    _check_output(result, expected)

    result = run_pip('freeze', '-f',
                     '%s/#egg=django-wikiapp' % checkout_path,
                     expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze -f %(repo)s/#egg=django-wikiapp
        -- stdout: --------------------
        -f %(repo)s/#egg=django-wikiapp
        ...-e %(repoC)s@...#egg=django_wikiapp-...
        ...""" % {'repoC': checkout_pathC, 'repo': checkout_path})
    _check_output(result, expected)

def test_freeze_with_local_option():
    """
    Test that wsgiref (from global site-packages) is reported normally, but not with --local.

    """
    reset_env()
    result = run_pip('install', 'initools==0.2')
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        INITools==0.2
        wsgiref==...
        <BLANKLINE>""")

    # The following check is broken (see
    # http://bitbucket.org/ianb/pip/issue/110).  For now we are simply
    # neutering this test, but if we can't find a way to fix it,
    # this whole function should be removed.

    # _check_output(result, expected)

    result = run_pip('freeze', '--local', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze --local
        -- stdout: --------------------
        INITools==0.2
        <BLANKLINE>""")
    _check_output(result, expected)

def test_freeze_with_requirement_option():
    """
    Test that new requirements are created correctly with --requirement hints

    """
    reset_env()
    ignores = textwrap.dedent("""\
        # Unchanged requirements below this line
        -r ignore.txt
        --requirement ignore.txt
        -Z ignore
        --always-unzip ignore
        -f http://ignore
        -i http://ignore
        --extra-index-url http://ignore
        --find-links http://ignore
        --index-url http://ignore
        """)
    write_file('hint.txt', textwrap.dedent("""\
        INITools==0.1
        NoExist==4.2
        """) + ignores)
    result = run_pip('install', 'initools==0.2')
    result = pip_install_local('simple')
    result = run_pip('freeze', '--requirement', 'hint.txt', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: pip freeze --requirement hint.txt
        -- stderr: --------------------
        Requirement file contains NoExist==4.2, but that package is not installed

        -- stdout: --------------------
        INITools==0.2
        """) + ignores + "## The following requirements were added by pip --freeze:..."
    _check_output(result, expected)
<commit_msg>Use PyPA's own test package instead of my personal repo (which I deleted due to move to Github).<commit_after>import sys
import re
import textwrap
from doctest import OutputChecker, ELLIPSIS
from tests.lib import reset_env, run_pip, write_file, get_env, pyversion, pip_install_local
from tests.lib.local_repos import local_checkout, local_repo

distribute_re = re.compile('^distribute==[0-9.]+\n', re.MULTILINE)

def _check_output(result, expected):
    checker = OutputChecker()
    actual = str(result)

    ## FIXME!  The following is a TOTAL hack.  For some reason the
    ## __str__ result for pkg_resources.Requirement gets downcased on
    ## Windows.  Since INITools is the only package we're installing
    ## in this file with funky case requirements, I'm forcibly
    ## upcasing it.  You can also normalize everything to lowercase,
    ## but then you have to remember to upcase <BLANKLINE>.  The right
    ## thing to do in the end is probably to find out how to report
    ## the proper fully-cased package name in our error message.
    if sys.platform == 'win32':
        actual = actual.replace('initools', 'INITools')

    # This allows our existing tests to work when run in a context
    # with distribute installed.
    actual = distribute_re.sub('', actual)

    def banner(msg):
        return '\n========== %s ==========\n' % msg
    assert checker.check_output(expected, actual, ELLIPSIS), banner('EXPECTED')+expected+banner('ACTUAL')+actual+banner(6*'=')

def test_freeze_basic():
    """
    Some tests of freeze, first we have to install some stuff.  Note that
    the test is a little crude at the end because Python 2.5+ adds egg
    info to the standard library, so stuff like wsgiref will show up in
    the freezing.  (Probably that should be accounted for in pip, but
    currently it is not).

    """
    env = reset_env()
    write_file('initools-req.txt', textwrap.dedent("""\
        simple==2.0
        # and something else to test out:
        simple2<=3.0
        """))
    result = pip_install_local('-r', env.scratch_path/'initools-req.txt')
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: pip freeze
        -- stdout: --------------------
        ...simple==2.0
        simple2==3.0...
        <BLANKLINE>""")
    _check_output(result, expected)

def test_freeze_svn():
    """Test freezing a svn checkout"""

    checkout_path = local_checkout('svn+http://svn.colorstudy.com/INITools/trunk')
    #svn internally stores windows drives as uppercase; we'll match that.
    checkout_path = checkout_path.replace('c:', 'C:')

    env = reset_env()
    result = env.run('svn', 'co', '-r10',
                     local_repo('svn+http://svn.colorstudy.com/INITools/trunk'),
                     'initools-trunk')
    result = env.run('python', 'setup.py', 'develop',
            cwd=env.scratch_path/ 'initools-trunk', expect_stderr=True)
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        -e %s@10#egg=INITools-0.3.1dev...-dev_r10
        ...""" % checkout_path)
    _check_output(result, expected)

def test_freeze_git_clone():
    """
    Test freezing a Git clone.

    """
    env = reset_env()
    result = env.run('git', 'clone', local_repo('git+http://github.com/pypa/pip-test-package.git'), 'pip-test-package')
    result = env.run('git', 'checkout', '7d654e66c8fa7149c165ddeffa5b56bc06619458',
            cwd=env.scratch_path / 'pip-test-package', expect_stderr=True)
    result = env.run('python', 'setup.py', 'develop',
            cwd=env.scratch_path / 'pip-test-package')
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        ...-e %s@...#egg=pip_test_package-...
        ...""" % local_checkout('git+http://github.com/pypa/pip-test-package.git'))
    _check_output(result, expected)

    result = run_pip('freeze', '-f',
                     '%s#egg=pip_test_package' % local_checkout('git+http://github.com/pypa/pip-test-package.git'),
                     expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: pip freeze -f %(repo)s#egg=pip_test_package
        -- stdout: --------------------
        -f %(repo)s#egg=pip_test_package...
        -e %(repo)s@...#egg=pip_test_package-0.1.1
        ...""" % {'repo': local_checkout('git+http://github.com/pypa/pip-test-package.git')})
    _check_output(result, expected)

def test_freeze_mercurial_clone():
    """
    Test freezing a Mercurial clone.

    """
    reset_env()
    env = get_env()
    result = env.run('hg', 'clone',
                     '-r', 'c9963c111e7c',
                     local_repo('hg+http://bitbucket.org/pypa/pip-test-package'),
                     'pip-test-package')
    result = env.run('python', 'setup.py', 'develop',
            cwd=env.scratch_path/'pip-test-package', expect_stderr=True)
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        ...-e %s@...#egg=pip_test_package-...
        ...""" % local_checkout('hg+http://bitbucket.org/pypa/pip-test-package'))
    _check_output(result, expected)

    result = run_pip('freeze', '-f',
                     '%s#egg=pip_test_package' % local_checkout('hg+http://bitbucket.org/pypa/pip-test-package'),
                     expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze -f %(repo)s#egg=pip_test_package
        -- stdout: --------------------
        -f %(repo)s#egg=pip_test_package
        ...-e %(repo)s@...#egg=pip_test_package-dev
        ...""" % {'repo': local_checkout('hg+http://bitbucket.org/pypa/pip-test-package')})
    _check_output(result, expected)

def test_freeze_bazaar_clone():
    """
    Test freezing a Bazaar clone.

    """

    checkout_path = local_checkout('bzr+http://bazaar.launchpad.net/%7Edjango-wikiapp/django-wikiapp/release-0.1')
    #bzr internally stores windows drives as uppercase; we'll match that.
    checkout_pathC = checkout_path.replace('c:', 'C:')

    reset_env()
    env = get_env()
    result = env.run('bzr', 'checkout', '-r', '174',
                     local_repo('bzr+http://bazaar.launchpad.net/%7Edjango-wikiapp/django-wikiapp/release-0.1'),
                     'django-wikiapp')
    result = env.run('python', 'setup.py', 'develop',
            cwd=env.scratch_path/'django-wikiapp')
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        ...-e %s@...#egg=django_wikiapp-...
        ...""" % checkout_pathC)
    _check_output(result, expected)

    result = run_pip('freeze', '-f',
                     '%s/#egg=django-wikiapp' % checkout_path,
                     expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze -f %(repo)s/#egg=django-wikiapp
        -- stdout: --------------------
        -f %(repo)s/#egg=django-wikiapp
        ...-e %(repoC)s@...#egg=django_wikiapp-...
        ...""" % {'repoC': checkout_pathC, 'repo': checkout_path})
    _check_output(result, expected)

def test_freeze_with_local_option():
    """
    Test that wsgiref (from global site-packages) is reported normally, but not with --local.

    """
    reset_env()
    result = run_pip('install', 'initools==0.2')
    result = run_pip('freeze', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze
        -- stdout: --------------------
        INITools==0.2
        wsgiref==...
        <BLANKLINE>""")

    # The following check is broken (see
    # http://bitbucket.org/ianb/pip/issue/110).  For now we are simply
    # neutering this test, but if we can't find a way to fix it,
    # this whole function should be removed.

    # _check_output(result, expected)

    result = run_pip('freeze', '--local', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: ...pip freeze --local
        -- stdout: --------------------
        INITools==0.2
        <BLANKLINE>""")
    _check_output(result, expected)

def test_freeze_with_requirement_option():
    """
    Test that new requirements are created correctly with --requirement hints

    """
    reset_env()
    ignores = textwrap.dedent("""\
        # Unchanged requirements below this line
        -r ignore.txt
        --requirement ignore.txt
        -Z ignore
        --always-unzip ignore
        -f http://ignore
        -i http://ignore
        --extra-index-url http://ignore
        --find-links http://ignore
        --index-url http://ignore
        """)
    write_file('hint.txt', textwrap.dedent("""\
        INITools==0.1
        NoExist==4.2
        """) + ignores)
    result = run_pip('install', 'initools==0.2')
    result = pip_install_local('simple')
    result = run_pip('freeze', '--requirement', 'hint.txt', expect_stderr=True)
    expected = textwrap.dedent("""\
        Script result: pip freeze --requirement hint.txt
        -- stderr: --------------------
        Requirement file contains NoExist==4.2, but that package is not installed

        -- stdout: --------------------
        INITools==0.2
        """) + ignores + "## The following requirements were added by pip --freeze:..."
    _check_output(result, expected)
<|endoftext|>


<commit_before>import pytest

@pytest.fixture(scope="module")
def learners():
	# add some learners that aren't trees
	from sklearn.tree import DecisionTreeRegressor
	return [
		DecisionTreeRegressor(criterion='friedman_mse', max_depth=5),
		DecisionTreeRegressor(criterion='friedman_mse', max_depth=3)
	]

class TestRegDistns():

	@pytest.fixture(scope="class")
	def reg_dists(self):
		# try importing these in the class but outside the fn
		from ngboost.distns import RegressionDistn
		from ngboost.distns import Normal, LogNormal, Exponential
		from ngboost.scores import LogScore, CRPScore
		return {
			Normal: [LogScore, CRPScore],
			LogNormal: [LogScore], #, CRPScore],
			Exponential: [LogScore, CRPScore]
			}

	@pytest.fixture(scope="class")
	def reg_data(self):
		from sklearn.datasets import load_boston
		from sklearn.model_selection import train_test_split
		X, Y = load_boston(True)
		return train_test_split(X, Y, test_size=0.2)

	def test_dists(self, learners, reg_dists, reg_data):
		from ngboost import NGBRegressor
		X_reg_train, X_reg_test, Y_reg_train, Y_reg_test = reg_data
		for Dist, Scores in reg_dists.items():
			for Score in Scores:
				for Learner in learners:
					# test early stopping features
					ngb = NGBRegressor(Dist=Dist, Score=Score, Base=Learner, verbose=False)
					ngb.fit(X_reg_train, Y_reg_train)
					y_pred = ngb.predict(X_reg_test)
					y_dist = ngb.pred_dist(X_reg_test)
					# test properties of output

	# test what happens when a dist that's not regression is passed in

# test survival stuff

class TestClsDistns():

	@pytest.fixture(scope="class")
	def cls_data(self):
		from sklearn.datasets import load_breast_cancer
		from sklearn.model_selection import train_test_split
		X, Y = load_breast_cancer(True)
		X_cls_train, X_cls_test, Y_cls_train, Y_cls_test  = train_test_split(X, Y, test_size=0.2)

	def test_bernoulli(self, learners, cls_data):
		from ngboost import NGBClassifier
		from ngboost.distns import Bernoulli
		from ngboost.scores import LogScore

		X_cls_train, X_cls_test, Y_cls_train, Y_cls_test = cls_data
		for Learner in learners:
			# test early stopping features
			# test other args, n_trees, LR, minibatching- args as fixture
			ngb = NGBClassifier(Dist=Bernoulli, Score=LogScore, Base=Learner, verbose=False)
			ngb.fit(X_cls_train, Y_cls_train)
			y_pred = ngb.predict(X_cls_test)
			y_prob = ngb.predict_proba(X_cls_test)
			y_dist = ngb.pred_dist(X_cls_test)
			# test properties of output

	def test_categorical(self, learners, cls_data):
		from ngboost import NGBClassifier
		from ngboost.distns import k_categorical
		from ngboost.scores import LogScore
		import numpy as np

		X_cls_train, X_cls_test, Y_cls_train, Y_cls_test = cls_data
		for K in [2,4,7]:
			Dist = k_categorical(K)
			Y_cls_train = np.random.randint(0,K,(len(Y_cls_train)))

			for Learner in learners:
				# test early stopping features
				ngb = NGBClassifier(Dist=Dist, Score=LogScore, Base=Learner, verbose=False)
				ngb.fit(X_cls_train, Y_cls_train)
				y_pred = ngb.predict(X_cls_test)
				y_prob = ngb.predict_proba(X_cls_test)
				y_dist = ngb.pred_dist(X_cls_test)
				# test properties of output

# test slicing and ._params<commit_msg>Update test_distns.py<commit_after>import pytest

@pytest.fixture(scope="module")
def learners():
	# add some learners that aren't trees
	from sklearn.tree import DecisionTreeRegressor
	return [
		DecisionTreeRegressor(criterion='friedman_mse', max_depth=5),
		DecisionTreeRegressor(criterion='friedman_mse', max_depth=3)
	]

class TestRegDistns():

	@pytest.fixture(scope="class")
	def reg_dists(self):
		# try importing these in the class but outside the fn
		from ngboost.distns import RegressionDistn
		from ngboost.distns import Normal, LogNormal, Exponential
		from ngboost.scores import LogScore, CRPScore
		return {
			Normal: [LogScore, CRPScore],
			LogNormal: [LogScore], #, CRPScore],
			Exponential: [LogScore, CRPScore]
			}

	@pytest.fixture(scope="class")
	def reg_data(self):
		from sklearn.datasets import load_boston
		from sklearn.model_selection import train_test_split
		X, Y = load_boston(True)
		return train_test_split(X, Y, test_size=0.2)

	def test_dists(self, learners, reg_dists, reg_data):
		from ngboost import NGBRegressor
		X_reg_train, X_reg_test, Y_reg_train, Y_reg_test = reg_data
		for Dist, Scores in reg_dists.items():
			for Score in Scores:
				for Learner in learners:
					# test early stopping features
					ngb = NGBRegressor(Dist=Dist, Score=Score, Base=Learner, verbose=False)
					ngb.fit(X_reg_train, Y_reg_train)
					y_pred = ngb.predict(X_reg_test)
					y_dist = ngb.pred_dist(X_reg_test)
					# test properties of output

	# test what happens when a dist that's not regression is passed in

# test survival stuff

class TestClsDistns():

	@pytest.fixture(scope="class")
	def cls_data(self):
		from sklearn.datasets import load_breast_cancer
		from sklearn.model_selection import train_test_split
		X, Y = load_breast_cancer(True)
		return train_test_split(X, Y, test_size=0.2)

	def test_bernoulli(self, learners, cls_data):
		from ngboost import NGBClassifier
		from ngboost.distns import Bernoulli
		from ngboost.scores import LogScore

		X_cls_train, X_cls_test, Y_cls_train, Y_cls_test = cls_data
		for Learner in learners:
			# test early stopping features
			# test other args, n_trees, LR, minibatching- args as fixture
			ngb = NGBClassifier(Dist=Bernoulli, Score=LogScore, Base=Learner, verbose=False)
			ngb.fit(X_cls_train, Y_cls_train)
			y_pred = ngb.predict(X_cls_test)
			y_prob = ngb.predict_proba(X_cls_test)
			y_dist = ngb.pred_dist(X_cls_test)
			# test properties of output

	def test_categorical(self, learners, cls_data):
		from ngboost import NGBClassifier
		from ngboost.distns import k_categorical
		from ngboost.scores import LogScore
		import numpy as np

		X_cls_train, X_cls_test, Y_cls_train, Y_cls_test = cls_data
		for K in [2,4,7]:
			Dist = k_categorical(K)
			Y_cls_train = np.random.randint(0,K,(len(Y_cls_train)))

			for Learner in learners:
				# test early stopping features
				ngb = NGBClassifier(Dist=Dist, Score=LogScore, Base=Learner, verbose=False)
				ngb.fit(X_cls_train, Y_cls_train)
				y_pred = ngb.predict(X_cls_test)
				y_prob = ngb.predict_proba(X_cls_test)
				y_dist = ngb.pred_dist(X_cls_test)
				# test properties of output

# test slicing and ._params<|endoftext|>


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

import os
import unittest

import pyrax

# This file needs to contain the actual credentials for a
# valid Rackspace Cloud account.
creds_file = os.path.expanduser("~/.rackspace_cloud_credentials")

class TestCase(unittest.TestCase):
    def setUp(self):
        pyrax.set_credential_file(creds_file)

    def tearDown(self):
        pyrax.clear_credentials()

    def test_cloudservers_images(self):
        imgs = pyrax.cloudservers.images.list()
        self.assert_(isinstance(imgs, list))

    def test_cloudfiles_base_container(self):
        conts = pyrax.cloudfiles.get_all_containers()
        self.assert_(isinstance(conts, list))

    def test_cloud_loadbalancers(self):
        lbs = pyrax.cloud_loadbalancers.list()
        self.assert_(isinstance(lbs, list))

    def test_cloud_db(self):
        flavors = pyrax.clouddatabases.list_flavors()
        self.assert_(isinstance(flavors, list))

if __name__ == "__main__":
    unittest.main()
<commit_msg>Fixed a typo in the test.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import unittest

import pyrax

# This file needs to contain the actual credentials for a
# valid Rackspace Cloud account.
creds_file = os.path.expanduser("~/.rackspace_cloud_credentials")

class TestCase(unittest.TestCase):
    def setUp(self):
        pyrax.set_credential_file(creds_file)

    def tearDown(self):
        pyrax.clear_credentials()

    def test_cloudservers_images(self):
        imgs = pyrax.cloudservers.images.list()
        self.assert_(isinstance(imgs, list))

    def test_cloudfiles_base_container(self):
        conts = pyrax.cloudfiles.get_all_containers()
        self.assert_(isinstance(conts, list))

    def test_cloud_loadbalancers(self):
        lbs = pyrax.cloud_loadbalancers.list()
        self.assert_(isinstance(lbs, list))

    def test_cloud_db(self):
        flavors = pyrax.cloud_databases.list_flavors()
        self.assert_(isinstance(flavors, list))

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


<commit_before>from __future__ import absolute_import

from base64 import b64encode
from collections import namedtuple
import csv
from datetime import date, datetime, timedelta
import io
import logging
import os
import uuid

from celery.schedules import crontab
from celery.task import periodic_task, task
from dateutil.relativedelta import relativedelta
from django.conf import settings
from django.db import Error, IntegrityError, connections
from django.db.models import F
import pytz

from corehq.apps.locations.models import SQLLocation
from corehq.apps.settings.views import get_qrcode
from corehq.apps.userreports.models import get_datasource_config
from corehq.apps.userreports.util import get_indicator_adapter
from corehq.const import SERVER_DATE_FORMAT
from corehq.form_processor.change_publishers import publish_case_saved
from corehq.form_processor.interfaces.dbaccessors import CaseAccessors
from corehq.sql_db.connections import get_icds_ucr_db_alias
from corehq.util.decorators import serial_task
from corehq.util.log import send_HTML_email
from corehq.util.soft_assert import soft_assert
from corehq.util.view_utils import reverse
from custom.icds_reports.models import AggChildHealthMonthly
from custom.icds_reports.reports.issnip_monthly_register import ISSNIPMonthlyReport
from custom.icds_reports.utils import zip_folder, create_pdf_file
from dimagi.utils.chunked import chunked
from dimagi.utils.logging import notify_exception
from six.moves import range

celery_task_logger = logging.getLogger('celery.task')

UCRAggregationTask = namedtuple("UCRAggregationTask", ['type', 'date'])

DASHBOARD_TEAM_MEMBERS = ['jemord', 'lbagnoli', 'ssrikrishnan', 'mharrison']
_dashboard_team_soft_assert = soft_assert(to=[
    '{}@{}'.format(member_id, 'dimagi.com') for member_id in DASHBOARD_TEAM_MEMBERS
])

@periodic_task(run_every=crontab(minute=30, hour=23), acks_late=True, queue='background_queue')
def run_move_ucr_data_into_aggregation_tables_task(date=None):
    move_ucr_data_into_aggregation_tables.delay(date)

@serial_task('move-ucr-data-into-aggregate-tables', timeout=30 * 60, queue='background_queue')
def move_ucr_data_into_aggregation_tables(date=None, intervals=2):
    date = date or datetime.utcnow().date()
    monthly_dates = []

    first_day_of_month = date.replace(day=1)
    for interval in range(intervals - 1, 0, -1):
        # calculate the last day of the previous months to send to the aggregation script
        first_day_next_month = first_day_of_month - relativedelta(months=interval - 1)
        monthly_dates.append(first_day_next_month - relativedelta(days=1))

    monthly_dates.append(date)

    db_alias = get_icds_ucr_db_alias()
    if db_alias:
        with connections[db_alias].cursor() as cursor:
            _create_aggregate_functions(cursor)
            _update_aggregate_locations_tables(cursor)

        aggregation_tasks = []

        for monthly_date in monthly_dates:
            calculation_date = monthly_date.strftime('%Y-%m-%d')
            aggregation_tasks.append(UCRAggregationTask('monthly', calculation_date))

        aggregation_tasks.append(UCRAggregationTask('daily', date.strftime('%Y-%m-%d')))
        aggregate_tables.delay(aggregation_tasks[0], aggregation_tasks[1:])

def _create_aggregate_functions(cursor):
    try:
        path = os.path.join(os.path.dirname(__file__), 'migrations', 'sql_templates', 'create_functions.sql')
        celery_task_logger.info("Starting icds reports create_functions")
        with open(path, "r") as sql_file:
            sql_to_execute = sql_file.read()
            cursor.execute(sql_to_execute)
        celery_task_logger.info("Ended icds reports create_functions")
    except Exception:
        # This is likely due to a change in the UCR models or aggregation script which should be rare
        # First step would be to look through this error to find what function is causing the error
        # and look for recent changes in this folder.
        _dashboard_team_soft_assert(False, "Unexpected occurred while creating functions in dashboard aggregation")
        raise

def _update_aggregate_locations_tables(cursor):
    try:
        path = os.path.join(os.path.dirname(__file__), 'sql_templates', 'update_locations_table.sql')
        celery_task_logger.info("Starting icds reports update_location_tables")
        with open(path, "r") as sql_file:
            sql_to_execute = sql_file.read()
            cursor.execute(sql_to_execute)
        celery_task_logger.info("Ended icds reports update_location_tables_sql")
    except IntegrityError:
        # This has occurred when there's a location upload, but not all locations were updated.
        # Some more details are here https://github.com/dimagi/commcare-hq/pull/18839
        # It's usually fixed by rebuild the location UCR table and running this task again, but
        # that PR should fix that issue
        _dashboard_team_soft_assert(False, "Error occurred while aggregating locations")
        raise
    except Exception:
        # I'm not sure what this one will be
        _dashboard_team_soft_assert(
            False, "Unexpected occurred while aggregating locations in dashboard aggregation")
        raise

@task(queue='background_queue', bind=True, default_retry_delay=15 * 60, acks_late=True)
def aggregate_tables(self, current_task, future_tasks):
    aggregation_type = current_task.type
    aggregation_date = current_task.date

    if aggregation_type == 'monthly':
        path = os.path.join(os.path.dirname(__file__), 'sql_templates', 'update_monthly_aggregate_tables.sql')
    elif aggregation_type == 'daily':
        path = os.path.join(os.path.dirname(__file__), 'sql_templates', 'update_daily_aggregate_table.sql')
    else:
        raise ValueError("Invalid aggregation type {}".format(aggregation_type))

    db_alias = get_icds_ucr_db_alias()
    if db_alias:
        with connections[db_alias].cursor() as cursor:
            with open(path, "r") as sql_file:
                sql_to_execute = sql_file.read()
                celery_task_logger.info(
                    "Starting icds reports {} update_{}_aggregate_tables".format(
                        aggregation_date, aggregation_type
                    )
                )

                try:
                    cursor.execute(sql_to_execute, {"date": aggregation_date})
                except Error as exc:
                    _dashboard_team_soft_assert(
                        False,
                        "{} aggregation failed on {} for {}. This task will be retried in 15 minutes".format(
                            aggregation_type, settings.SERVER_ENVIRONMENT, aggregation_date
                        )
                    )
                    notify_exception(
                        None,
                        message="Error occurred during ICDS aggregation",
                        details={
                            'type': aggregation_type,
                            'date': aggregation_date,
                            'error': exc,
                        }
                    )
                    self.retry(exc=exc)

                celery_task_logger.info(
                    "Ended icds reports {} update_{}_aggregate_tables".format(
                        aggregation_date, aggregation_type
                    )
                )

    if future_tasks:
        aggregate_tables.delay(future_tasks[0], future_tasks[1:])
    else:
        # temporary soft assert to verify it's completing
        _dashboard_team_soft_assert(False, "Aggregation completed on {}".format(settings.SERVER_ENVIRONMENT))
        celery_task_logger.info("Aggregation has completed")

@periodic_task(
    queue='background_queue',
    run_every=crontab(day_of_week='tuesday,thursday,saturday', minute=0, hour=21),
    acks_late=True
)
def recalculate_stagnant_cases():
    domain = 'icds-cas'
    config_ids = [
        'static-icds-cas-static-ccs_record_cases_monthly_v2',
        'static-icds-cas-static-ccs_record_cases_monthly_tableau_v2',
        'static-icds-cas-static-child_cases_monthly_v2',
        'static-icds-cas-static-child_cases_monthly_tableau_v2',
    ]

    stagnant_cases = set()

    for config_id in config_ids:
        config, is_static = get_datasource_config(config_id, domain)
        adapter = get_indicator_adapter(config)
        case_ids = _find_stagnant_cases(adapter)
        celery_task_logger.info(
            "Found {} stagnant cases in config {}".format(len(case_ids), config_id)
        )
        stagnant_cases = stagnant_cases.union(set(case_ids))
        celery_task_logger.info(
            "Total number of stagant cases is now {}".format(len(stagnant_cases))
        )

    case_accessor = CaseAccessors(domain)
    num_stagnant_cases = len(stagnant_cases)
    current_case_num = 0
    for case_ids in chunked(stagnant_cases, 1000):
        current_case_num += len(case_ids)
        cases = case_accessor.get_cases(list(case_ids))
        for case in cases:
            publish_case_saved(case, send_post_save_signal=False)
        celery_task_logger.info(
            "Resaved {} / {} cases".format(current_case_num, num_stagnant_cases)
        )

def _find_stagnant_cases(adapter):
    stagnant_date = datetime.utcnow() - timedelta(days=26)
    table = adapter.get_table()
    query = adapter.get_query_object()
    query = query.with_entities(table.columns.doc_id).filter(
        table.columns.inserted_at <= stagnant_date
    ).distinct()
    return query.all()

india_timezone = pytz.timezone('Asia/Kolkata')

@task(queue='background_queue')
def prepare_issnip_monthly_register_reports(domain, user, awcs, pdf_format, month, year):

    utc_now = datetime.now(pytz.utc)
    india_now = utc_now.astimezone(india_timezone)

    selected_date = date(year, month, 1)
    report_context = {
        'reports': []
    }

    pdf_files = []

    for awc in awcs:
        pdf_hash = uuid.uuid4().hex

        awc_location = SQLLocation.objects.get(location_id=awc)
        pdf_files.append({
            'uuid': pdf_hash,
            'location_name': awc_location.name.replace(' ', '_')
        })
        report = ISSNIPMonthlyReport(config={
            'awc_id': awc_location.location_id,
            'month': selected_date,
            'domain': domain
        })
        qrcode = get_qrcode("{} {}".format(
            awc_location.site_code,
            india_now.strftime('%d %b %Y')
        ))
        context = {
            'qrcode_64': b64encode(qrcode),
            'report': report
        }

        if pdf_format == 'one':
            report_context['reports'].append(context)
        else:
            report_context['reports'] = [context]
            create_pdf_file(
                pdf_hash,
                report_context
            )

    if pdf_format == 'many':
        cache_key = zip_folder(pdf_files)
    else:
        cache_key = create_pdf_file(uuid.uuid4().hex, report_context)

    params = {
        'domain': domain,
        'uuid': cache_key,
        'format': pdf_format
    }

    return {
        'domain': domain,
        'uuid': cache_key,
        'format': pdf_format,
        'link': reverse('icds_download_pdf', params=params, absolute=True, kwargs={'domain': domain})
    }

@task(queue='background_queue')
def icds_data_validation(day):
    """Checks all AWCs to validate that there will be no inconsistencies in the
    reporting dashboard.
    """

    # agg tables store the month like YYYY-MM-01
    month = day
    month.replace(day=1)
    return_values = ('state_name', 'district_name', 'block_name', 'supervisor_name', 'awc_name')

    bad_wasting_awcs = AggChildHealthMonthly.objects.filter(
        month=month, aggregation_level=5
    ).exclude(
        weighed_and_height_measured_in_month=(
            F('wasting_moderate') + F('wasting_severe') + F('wasting_normal')
        )
    ).values_list(*return_values)

    bad_stunting_awcs = AggChildHealthMonthly.objects.filter(month=month, aggregation_level=5).exclude(
        height_measured_in_month=(
            F('stunting_severe') + F('stunting_moderate') + F('stunting_normal')
        )
    ).values_list(*return_values)

    bad_underweight_awcs = AggChildHealthMonthly.objects.filter(month=month, aggregation_level=5).exclude(
        nutrition_status_weighed=(
            F('nutrition_status_normal') +
            F('nutrition_status_moderately_underweight') +
            F('nutrition_status_severely_underweight')
        )
    ).values_list(*return_values)

    bad_lbw_awcs = AggChildHealthMonthly.objects.filter(month=month, aggregation_level=5).exclude(
        weighed_and_born_in_month__gt=(
            F('low_birth_weight_in_month')
        )
    ).values_list(*return_values)

    csv_file = io.BytesIO()
    writer = csv.writer(csv_file)
    writer.writerow(('type',) + return_values)
    _icds_add_awcs_to_file(writer, 'wasting', bad_wasting_awcs)
    _icds_add_awcs_to_file(writer, 'stunting', bad_stunting_awcs)
    _icds_add_awcs_to_file(writer, 'underweight', bad_underweight_awcs)
    _icds_add_awcs_to_file(writer, 'low_birth_weight', bad_lbw_awcs)

    email_content = """
    Incorrect wasting AWCs: {bad_wasting_awcs}
    Incorrect stunting AWCs: {bad_stunting_awcs}
    Incorrect underweight AWCs: {bad_underweight_awcs}
    Incorrect low birth weight AWCs: {bad_lbw_awcs}

    Please see attached file for more details
    """.format(
        bad_wasting_awcs=len(bad_wasting_awcs),
        bad_stunting_awcs=len(bad_stunting_awcs),
        bad_underweight_awcs=len(bad_underweight_awcs),
        bad_lbw_awcs=len(bad_lbw_awcs),
    )

    filename = month.strftime('validation_results_%s.csv' % SERVER_DATE_FORMAT)
    send_HTML_email(
        'ICDS Dashboard Validation Results', DASHBOARD_TEAM_MEMBERS, email_content,
        file_attachments=[{'file_obj': csv_file, 'title': filename, 'mimetype': 'text/csv'}],
    )

def _icds_add_awcs_to_file(csv_writer, error_type, rows):
    for row in rows:
        csv_writer.writerow((error_type, ) + row)
<commit_msg>Run validation after aggregation script completes.<commit_after>from __future__ import absolute_import

from base64 import b64encode
from collections import namedtuple
import csv
from datetime import date, datetime, timedelta
import io
import logging
import os
import uuid

from celery.schedules import crontab
from celery.task import periodic_task, task
from dateutil.relativedelta import relativedelta
from django.conf import settings
from django.db import Error, IntegrityError, connections
from django.db.models import F
import pytz

from corehq.apps.locations.models import SQLLocation
from corehq.apps.settings.views import get_qrcode
from corehq.apps.userreports.models import get_datasource_config
from corehq.apps.userreports.util import get_indicator_adapter
from corehq.const import SERVER_DATE_FORMAT
from corehq.form_processor.change_publishers import publish_case_saved
from corehq.form_processor.interfaces.dbaccessors import CaseAccessors
from corehq.sql_db.connections import get_icds_ucr_db_alias
from corehq.util.decorators import serial_task
from corehq.util.log import send_HTML_email
from corehq.util.soft_assert import soft_assert
from corehq.util.view_utils import reverse
from custom.icds_reports.models import AggChildHealthMonthly
from custom.icds_reports.reports.issnip_monthly_register import ISSNIPMonthlyReport
from custom.icds_reports.utils import zip_folder, create_pdf_file
from dimagi.utils.chunked import chunked
from dimagi.utils.logging import notify_exception
from six.moves import range

celery_task_logger = logging.getLogger('celery.task')

UCRAggregationTask = namedtuple("UCRAggregationTask", ['type', 'date'])

DASHBOARD_TEAM_MEMBERS = ['jemord', 'lbagnoli', 'ssrikrishnan', 'mharrison']
_dashboard_team_soft_assert = soft_assert(to=[
    '{}@{}'.format(member_id, 'dimagi.com') for member_id in DASHBOARD_TEAM_MEMBERS
])

@periodic_task(run_every=crontab(minute=30, hour=23), acks_late=True, queue='background_queue')
def run_move_ucr_data_into_aggregation_tables_task(date=None):
    move_ucr_data_into_aggregation_tables.delay(date)

@serial_task('move-ucr-data-into-aggregate-tables', timeout=30 * 60, queue='background_queue')
def move_ucr_data_into_aggregation_tables(date=None, intervals=2):
    date = date or datetime.utcnow().date()
    monthly_dates = []

    first_day_of_month = date.replace(day=1)
    for interval in range(intervals - 1, 0, -1):
        # calculate the last day of the previous months to send to the aggregation script
        first_day_next_month = first_day_of_month - relativedelta(months=interval - 1)
        monthly_dates.append(first_day_next_month - relativedelta(days=1))

    monthly_dates.append(date)

    db_alias = get_icds_ucr_db_alias()
    if db_alias:
        with connections[db_alias].cursor() as cursor:
            _create_aggregate_functions(cursor)
            _update_aggregate_locations_tables(cursor)

        aggregation_tasks = []

        for monthly_date in monthly_dates:
            calculation_date = monthly_date.strftime('%Y-%m-%d')
            aggregation_tasks.append(UCRAggregationTask('monthly', calculation_date))

        aggregation_tasks.append(UCRAggregationTask('daily', date.strftime('%Y-%m-%d')))
        aggregate_tables.delay(aggregation_tasks[0], aggregation_tasks[1:])

def _create_aggregate_functions(cursor):
    try:
        path = os.path.join(os.path.dirname(__file__), 'migrations', 'sql_templates', 'create_functions.sql')
        celery_task_logger.info("Starting icds reports create_functions")
        with open(path, "r") as sql_file:
            sql_to_execute = sql_file.read()
            cursor.execute(sql_to_execute)
        celery_task_logger.info("Ended icds reports create_functions")
    except Exception:
        # This is likely due to a change in the UCR models or aggregation script which should be rare
        # First step would be to look through this error to find what function is causing the error
        # and look for recent changes in this folder.
        _dashboard_team_soft_assert(False, "Unexpected occurred while creating functions in dashboard aggregation")
        raise

def _update_aggregate_locations_tables(cursor):
    try:
        path = os.path.join(os.path.dirname(__file__), 'sql_templates', 'update_locations_table.sql')
        celery_task_logger.info("Starting icds reports update_location_tables")
        with open(path, "r") as sql_file:
            sql_to_execute = sql_file.read()
            cursor.execute(sql_to_execute)
        celery_task_logger.info("Ended icds reports update_location_tables_sql")
    except IntegrityError:
        # This has occurred when there's a location upload, but not all locations were updated.
        # Some more details are here https://github.com/dimagi/commcare-hq/pull/18839
        # It's usually fixed by rebuild the location UCR table and running this task again, but
        # that PR should fix that issue
        _dashboard_team_soft_assert(False, "Error occurred while aggregating locations")
        raise
    except Exception:
        # I'm not sure what this one will be
        _dashboard_team_soft_assert(
            False, "Unexpected occurred while aggregating locations in dashboard aggregation")
        raise

@task(queue='background_queue', bind=True, default_retry_delay=15 * 60, acks_late=True)
def aggregate_tables(self, current_task, future_tasks):
    aggregation_type = current_task.type
    aggregation_date = current_task.date

    if aggregation_type == 'monthly':
        path = os.path.join(os.path.dirname(__file__), 'sql_templates', 'update_monthly_aggregate_tables.sql')
    elif aggregation_type == 'daily':
        path = os.path.join(os.path.dirname(__file__), 'sql_templates', 'update_daily_aggregate_table.sql')
    else:
        raise ValueError("Invalid aggregation type {}".format(aggregation_type))

    db_alias = get_icds_ucr_db_alias()
    if db_alias:
        with connections[db_alias].cursor() as cursor:
            with open(path, "r") as sql_file:
                sql_to_execute = sql_file.read()
                celery_task_logger.info(
                    "Starting icds reports {} update_{}_aggregate_tables".format(
                        aggregation_date, aggregation_type
                    )
                )

                try:
                    cursor.execute(sql_to_execute, {"date": aggregation_date})
                except Error as exc:
                    _dashboard_team_soft_assert(
                        False,
                        "{} aggregation failed on {} for {}. This task will be retried in 15 minutes".format(
                            aggregation_type, settings.SERVER_ENVIRONMENT, aggregation_date
                        )
                    )
                    notify_exception(
                        None,
                        message="Error occurred during ICDS aggregation",
                        details={
                            'type': aggregation_type,
                            'date': aggregation_date,
                            'error': exc,
                        }
                    )
                    self.retry(exc=exc)

                celery_task_logger.info(
                    "Ended icds reports {} update_{}_aggregate_tables".format(
                        aggregation_date, aggregation_type
                    )
                )

    if future_tasks:
        aggregate_tables.delay(future_tasks[0], future_tasks[1:])
    else:
        # temporary soft assert to verify it's completing
        _dashboard_team_soft_assert(False, "Aggregation completed on {}".format(settings.SERVER_ENVIRONMENT))
        celery_task_logger.info("Aggregation has completed")
        icds_data_validation.delay(aggregation_date)

@periodic_task(
    queue='background_queue',
    run_every=crontab(day_of_week='tuesday,thursday,saturday', minute=0, hour=21),
    acks_late=True
)
def recalculate_stagnant_cases():
    domain = 'icds-cas'
    config_ids = [
        'static-icds-cas-static-ccs_record_cases_monthly_v2',
        'static-icds-cas-static-ccs_record_cases_monthly_tableau_v2',
        'static-icds-cas-static-child_cases_monthly_v2',
        'static-icds-cas-static-child_cases_monthly_tableau_v2',
    ]

    stagnant_cases = set()

    for config_id in config_ids:
        config, is_static = get_datasource_config(config_id, domain)
        adapter = get_indicator_adapter(config)
        case_ids = _find_stagnant_cases(adapter)
        celery_task_logger.info(
            "Found {} stagnant cases in config {}".format(len(case_ids), config_id)
        )
        stagnant_cases = stagnant_cases.union(set(case_ids))
        celery_task_logger.info(
            "Total number of stagant cases is now {}".format(len(stagnant_cases))
        )

    case_accessor = CaseAccessors(domain)
    num_stagnant_cases = len(stagnant_cases)
    current_case_num = 0
    for case_ids in chunked(stagnant_cases, 1000):
        current_case_num += len(case_ids)
        cases = case_accessor.get_cases(list(case_ids))
        for case in cases:
            publish_case_saved(case, send_post_save_signal=False)
        celery_task_logger.info(
            "Resaved {} / {} cases".format(current_case_num, num_stagnant_cases)
        )

def _find_stagnant_cases(adapter):
    stagnant_date = datetime.utcnow() - timedelta(days=26)
    table = adapter.get_table()
    query = adapter.get_query_object()
    query = query.with_entities(table.columns.doc_id).filter(
        table.columns.inserted_at <= stagnant_date
    ).distinct()
    return query.all()

india_timezone = pytz.timezone('Asia/Kolkata')

@task(queue='background_queue')
def prepare_issnip_monthly_register_reports(domain, user, awcs, pdf_format, month, year):

    utc_now = datetime.now(pytz.utc)
    india_now = utc_now.astimezone(india_timezone)

    selected_date = date(year, month, 1)
    report_context = {
        'reports': []
    }

    pdf_files = []

    for awc in awcs:
        pdf_hash = uuid.uuid4().hex

        awc_location = SQLLocation.objects.get(location_id=awc)
        pdf_files.append({
            'uuid': pdf_hash,
            'location_name': awc_location.name.replace(' ', '_')
        })
        report = ISSNIPMonthlyReport(config={
            'awc_id': awc_location.location_id,
            'month': selected_date,
            'domain': domain
        })
        qrcode = get_qrcode("{} {}".format(
            awc_location.site_code,
            india_now.strftime('%d %b %Y')
        ))
        context = {
            'qrcode_64': b64encode(qrcode),
            'report': report
        }

        if pdf_format == 'one':
            report_context['reports'].append(context)
        else:
            report_context['reports'] = [context]
            create_pdf_file(
                pdf_hash,
                report_context
            )

    if pdf_format == 'many':
        cache_key = zip_folder(pdf_files)
    else:
        cache_key = create_pdf_file(uuid.uuid4().hex, report_context)

    params = {
        'domain': domain,
        'uuid': cache_key,
        'format': pdf_format
    }

    return {
        'domain': domain,
        'uuid': cache_key,
        'format': pdf_format,
        'link': reverse('icds_download_pdf', params=params, absolute=True, kwargs={'domain': domain})
    }

@task(queue='background_queue')
def icds_data_validation(day):
    """Checks all AWCs to validate that there will be no inconsistencies in the
    reporting dashboard.
    """

    # agg tables store the month like YYYY-MM-01
    month = day
    month.replace(day=1)
    return_values = ('state_name', 'district_name', 'block_name', 'supervisor_name', 'awc_name')

    bad_wasting_awcs = AggChildHealthMonthly.objects.filter(
        month=month, aggregation_level=5
    ).exclude(
        weighed_and_height_measured_in_month=(
            F('wasting_moderate') + F('wasting_severe') + F('wasting_normal')
        )
    ).values_list(*return_values)

    bad_stunting_awcs = AggChildHealthMonthly.objects.filter(month=month, aggregation_level=5).exclude(
        height_measured_in_month=(
            F('stunting_severe') + F('stunting_moderate') + F('stunting_normal')
        )
    ).values_list(*return_values)

    bad_underweight_awcs = AggChildHealthMonthly.objects.filter(month=month, aggregation_level=5).exclude(
        nutrition_status_weighed=(
            F('nutrition_status_normal') +
            F('nutrition_status_moderately_underweight') +
            F('nutrition_status_severely_underweight')
        )
    ).values_list(*return_values)

    bad_lbw_awcs = AggChildHealthMonthly.objects.filter(month=month, aggregation_level=5).exclude(
        weighed_and_born_in_month__gt=(
            F('low_birth_weight_in_month')
        )
    ).values_list(*return_values)

    csv_file = io.BytesIO()
    writer = csv.writer(csv_file)
    writer.writerow(('type',) + return_values)
    _icds_add_awcs_to_file(writer, 'wasting', bad_wasting_awcs)
    _icds_add_awcs_to_file(writer, 'stunting', bad_stunting_awcs)
    _icds_add_awcs_to_file(writer, 'underweight', bad_underweight_awcs)
    _icds_add_awcs_to_file(writer, 'low_birth_weight', bad_lbw_awcs)

    email_content = """
    Incorrect wasting AWCs: {bad_wasting_awcs}
    Incorrect stunting AWCs: {bad_stunting_awcs}
    Incorrect underweight AWCs: {bad_underweight_awcs}
    Incorrect low birth weight AWCs: {bad_lbw_awcs}

    Please see attached file for more details
    """.format(
        bad_wasting_awcs=len(bad_wasting_awcs),
        bad_stunting_awcs=len(bad_stunting_awcs),
        bad_underweight_awcs=len(bad_underweight_awcs),
        bad_lbw_awcs=len(bad_lbw_awcs),
    )

    filename = month.strftime('validation_results_%s.csv' % SERVER_DATE_FORMAT)
    send_HTML_email(
        'ICDS Dashboard Validation Results', DASHBOARD_TEAM_MEMBERS, email_content,
        file_attachments=[{'file_obj': csv_file, 'title': filename, 'mimetype': 'text/csv'}],
    )

def _icds_add_awcs_to_file(csv_writer, error_type, rows):
    for row in rows:
        csv_writer.writerow((error_type, ) + row)
<|endoftext|>


<commit_before>import json
import os

from datetime import datetime, timedelta

from corehq.apps.locations.models import SQLLocation
from corehq.apps.reports_core.filters import Choice
from corehq.apps.userreports.models import ReportConfiguration
from corehq.apps.userreports.reports.factory import ReportFactory
from corehq.util.couch import get_document_or_not_found
from dimagi.utils.dates import DateSpan

class MPRData(object):
    resource_file = 'resources/block_mpr.json'

class ASRData(object):
    resource_file = 'resources/block_asr.json'

class ICDSData(object):

    def __init__(self, domain, filters, report_id):
        report_config = ReportFactory.from_spec(
            get_document_or_not_found(
                ReportConfiguration,
                domain,
                report_id
            )
        )
        report_config.set_filter_values(filters)
        self.report_config = report_config

    def data(self):
        return self.report_config.get_data()

class ICDSMixin(object):
    has_sections = False
    posttitle = None

    def __init__(self, config):
        self.config = config

    @property
    def subtitle(self):
        return []

    @property
    def headers(self):
        return []

    @property
    def rows(self):
        return [[]]

    @property
    def sources(self):
        with open(os.path.join(os.path.dirname(__file__), self.resource_file)) as f:
            return json.loads(f.read())[self.slug]

    @property
    def selected_location(self):
        if self.config['location_id']:
            return SQLLocation.objects.get(
                location_id=self.config['location_id']
            )

    @property
    def awc(self):
        if self.config['location_id']:
            return self.selected_location.get_descendants(include_self=True).filter(
                location_type__name='awc'
            )

    @property
    def awc_number(self):
        if self.awc:
            return len(
                [
                    loc for loc in self.awc
                    if 'test' not in loc.metadata and loc.metadata.get('test', '').lower() != 'yes'
                ]
            )

    def custom_data(self, selected_location, domain):
        data = {}

        for config in self.sources['data_source']:
            filters = {}
            if selected_location:
                key = selected_location.location_type.name.lower() + '_id'
                filters = {
                    key: [Choice(value=selected_location.location_id, display=selected_location.name)]
                }
            if 'date_filter_field' in config:
                filters.update({config['date_filter_field']: self.config['date_span']})
            if 'filter' in config:
                for fil in config['filter']:
                    if 'type' in fil:
                        now = datetime.now()
                        start_date = now if 'start' not in fil else now + timedelta(days=fil['start'])
                        end_date = now if 'end' not in fil else now + timedelta(days=fil['end'])
                        datespan = DateSpan(start_date, end_date)
                        filters.update({fil['column']: datespan})
                    else:
                        filters.update({
                            fil['column']: {
                                'operator': fil['operator'],
                                'operand': fil['value']
                            }
                        })

            report_data = ICDSData(domain, filters, config['id']).data()
            for column in config['columns']:
                column_agg_func = column['agg_fun']
                column_name = column['column_name']
                column_data = 0
                if column_agg_func == 'sum':
                    column_data = sum([x.get(column_name, 0) for x in report_data])
                elif column_agg_func == 'count':
                    column_data = len(report_data)
                elif column_agg_func == 'count_if':
                    value = column['condition']['value']
                    if isinstance(str, value):
                        statement = 'val.get("{0}", "") {1} "{2}"'
                    else:
                        statement = 'val.get("{0}", 0) {1} {2}'
                    condition = statement.format(
                        column_name,
                        column['condition']['operator'],
                        value
                    )
                    column_data = len([val for val in report_data if eval(condition)])
                elif column_agg_func == 'avg':
                    values = [x.get(column_name, 0) for x in report_data]
                    column_data = sum(values) / (len(values) or 1)
                column_display = column_name if 'column_in_report' not in column else column['column_in_report']
                data.update({
                    column_display: data.get(column_display, 0) + column_data
                })
        return data
<commit_msg>Typo in block reports<commit_after>import json
import os

from datetime import datetime, timedelta

from corehq.apps.locations.models import SQLLocation
from corehq.apps.reports_core.filters import Choice
from corehq.apps.userreports.models import ReportConfiguration
from corehq.apps.userreports.reports.factory import ReportFactory
from corehq.util.couch import get_document_or_not_found
from dimagi.utils.dates import DateSpan

class MPRData(object):
    resource_file = 'resources/block_mpr.json'

class ASRData(object):
    resource_file = 'resources/block_asr.json'

class ICDSData(object):

    def __init__(self, domain, filters, report_id):
        report_config = ReportFactory.from_spec(
            get_document_or_not_found(
                ReportConfiguration,
                domain,
                report_id
            )
        )
        report_config.set_filter_values(filters)
        self.report_config = report_config

    def data(self):
        return self.report_config.get_data()

class ICDSMixin(object):
    has_sections = False
    posttitle = None

    def __init__(self, config):
        self.config = config

    @property
    def subtitle(self):
        return []

    @property
    def headers(self):
        return []

    @property
    def rows(self):
        return [[]]

    @property
    def sources(self):
        with open(os.path.join(os.path.dirname(__file__), self.resource_file)) as f:
            return json.loads(f.read())[self.slug]

    @property
    def selected_location(self):
        if self.config['location_id']:
            return SQLLocation.objects.get(
                location_id=self.config['location_id']
            )

    @property
    def awc(self):
        if self.config['location_id']:
            return self.selected_location.get_descendants(include_self=True).filter(
                location_type__name='awc'
            )

    @property
    def awc_number(self):
        if self.awc:
            return len(
                [
                    loc for loc in self.awc
                    if 'test' not in loc.metadata and loc.metadata.get('test', '').lower() != 'yes'
                ]
            )

    def custom_data(self, selected_location, domain):
        data = {}

        for config in self.sources['data_source']:
            filters = {}
            if selected_location:
                key = selected_location.location_type.name.lower() + '_id'
                filters = {
                    key: [Choice(value=selected_location.location_id, display=selected_location.name)]
                }
            if 'date_filter_field' in config:
                filters.update({config['date_filter_field']: self.config['date_span']})
            if 'filter' in config:
                for fil in config['filter']:
                    if 'type' in fil:
                        now = datetime.now()
                        start_date = now if 'start' not in fil else now + timedelta(days=fil['start'])
                        end_date = now if 'end' not in fil else now + timedelta(days=fil['end'])
                        datespan = DateSpan(start_date, end_date)
                        filters.update({fil['column']: datespan})
                    else:
                        filters.update({
                            fil['column']: {
                                'operator': fil['operator'],
                                'operand': fil['value']
                            }
                        })

            report_data = ICDSData(domain, filters, config['id']).data()
            for column in config['columns']:
                column_agg_func = column['agg_fun']
                column_name = column['column_name']
                column_data = 0
                if column_agg_func == 'sum':
                    column_data = sum([x.get(column_name, 0) for x in report_data])
                elif column_agg_func == 'count':
                    column_data = len(report_data)
                elif column_agg_func == 'count_if':
                    value = column['condition']['value']
                    if isinstance(value, str):
                        statement = 'val.get("{0}", "") {1} "{2}"'
                    else:
                        statement = 'val.get("{0}", 0) {1} {2}'
                    condition = statement.format(
                        column_name,
                        column['condition']['operator'],
                        value
                    )
                    column_data = len([val for val in report_data if eval(condition)])
                elif column_agg_func == 'avg':
                    values = [x.get(column_name, 0) for x in report_data]
                    column_data = sum(values) / (len(values) or 1)
                column_display = column_name if 'column_in_report' not in column else column['column_in_report']
                data.update({
                    column_display: data.get(column_display, 0) + column_data
                })
        return data
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
"""
Import from s9y RSS export

Warning: this is specifically written for

https://starship-factory.ch/rss.php?version=2.0&all=1

and is pretty much a one-off script. Don't expect it to work on
other RSS feeds, it might not work or break your site/eat your babies/...

"""
import os
import re
import sys
import json
import time
import rfc822
import hashlib
import urlparse
import datetime
from contextlib import contextmanager
import requests
from lxml import objectify, html, etree
from lxml.html.clean import clean_html

from django.db import transaction
from django.utils.html import escape
from django.utils.text import slugify
from django.utils.timezone import make_aware, utc
from django.core.files.images import ImageFile
from django.contrib.sites.models import Site
from django.contrib.redirects.models import Redirect

from simplecms.models import TextSnippet, ImageSnippet, Content
from zinnia.models.entry import Entry
from zinnia.models.category import Category

_url_cache = {}

def my_slugify(value):
    """
    slugify with umlauts
    """
    translation_table = {
        u'ß': 'ss',
        u'ä': 'ae',
        u'ö': 'oe',
        u'ü': 'ue',
        u'é': 'e',
        u'Ä': 'Ae',
        u'Ö': 'Oe',
        u'Ü': 'Ue',
        u'É': 'E',
    }
    for k, v in translation_table.items():
        value = value.replace(k, v)
    return slugify(value)

def follow_redirects(url, verify=True):
    """
    follow any redirects and return the resulting url
    """
    s = requests.Session()
    s.max_redirects = 10  # FIXME: magic number...
    s.verify = verify
    try:
        resp = s.head(url, allow_redirects=True)
    except requests.exceptions.SSLError, e:
        if verify:
            return follow_redirects(url, verify=False)
        else:
            print >> sys.stderr, e, url
            return url
    except requests.exceptions.ConnectionError, e:
        print >> sys.stderr, e, url
        return url
    return resp.url

def transform_image_url(url):
    """
    rewrite the url to the image to one that is closer to the original size
    """
    from django.utils.http import urlencode
    url = url.replace("serendipityThumb.", "")
    if "admin_image_selector.php" in url:
        width_key = 'serendipity[resizeWidth]'
        height_key = 'serendipity[resizeHeight]'
        parsed_url = urlparse.urlparse(url)
        qs = parsed_url.query
        parsed_qs = urlparse.parse_qs(qs)

        if height_key in parsed_qs:
            parsed_qs[height_key] = "10000"
        if width_key in parsed_qs:
            parsed_qs[width_key] = "10000"

        qs = urlencode(parsed_qs, doseq=1)
        url = parsed_url._replace(query=qs).geturl()
    return url

def add_redirect(full_old_url, new_path):
    """
    add a redirect entry. for example
    full_old_url=http://example.org/blog/post1/
    new_path=/2014/05/06/post1/
    """
    current_site = Site.objects.get_current()
    parsed_old = urlparse.urlparse(full_old_url)
    if parsed_old.query:
        # bail out: can only handle paths
        return
    # theres a unique constraint on (site_id, old_path), so get rid of old redirects
    Redirect.objects.filter(site=current_site,
                            old_path=parsed_old.path).delete()
    Redirect.objects.get_or_create(site=current_site,
                                   old_path=parsed_old.path,
                                   new_path=new_path)

def image_cache_path(url):
    key = hashlib.md5(url).hexdigest()
    if not os.isdir("image_cache"):
        os.mkdir("image_cache")
    return os.path.join("image_cache", key)

def image_cache_get(url):
    """
    get a file handle for the image at ``url``
    """
    path = image_cache_path(url)
    if os.path.exists(path):
        return open(path)

@contextmanager
def image_cache_put(url):
    path = image_cache_path(url)
    fp = open(path, "wb")
    yield fp
    fp.close()

def fetch_image(url):
    url = transform_image_url(url)
    img = image_cache_get(url)
    if img:
        return img
    resp = requests.get(url, stream=True)
    if resp.status_code == 200:
        with image_cache_put(url) as fp:
            for chunk in resp.iter_content(1024):
                fp.write(chunk)
    return image_cache_get(url)

def resolve_url(url):
    if url in _url_cache:
        return _url_cache[url]
    dest_url = follow_redirects(url)
    _url_cache[url] = dest_url
    return dest_url

class HTMLWalker(object):
    """
    walk an element tree; when encountering an image, save the elements encountered
    from ``self.current_snippet`` into ``self.snippets`` and start a new snippet.
    """
    self_closing = ["img"]

    def __init__(self):
        self.current_snippet = []
        self.snippets = []
        self.current_type = "text"

    def add(self, blob):
        self.current_snippet.append(blob)

    def done(self, new_type="text"):
        self.snippets.append((self.current_type, self.current_snippet))
        self.current_type = new_type
        self.current_snippet = []

    def walk(self, elem, indent=0, stack=None):
        stack = stack or []
        space = indent * 4 * " "
        space = ""

        if elem.tag == "img":
            if "src" not in elem.attrib:
                return
            # close any open tags
            for e in list(reversed(stack))[1:]:
                self.add("%s</%s>" % (space, e.tag,))
            self.done(new_type="image")

        # special case: thumbnail linked to the full image
        if elem.tag == "a" and elem.attrib.get("class", "") == "serendipity_image_link":
            self.done(new_type="image")
            title = elem.find("img").attrib.get("title", "")
            self.add("" %
                     (elem.attrib.get("href"),
                      title, title))
            self.done()
            return

        if elem.tag == "div" and len(elem.xpath("//*")) == 0:
            return

        if elem.attrib:
            attrs = " " + " ".join('%s="%s"' % (escape(k), escape(v))
                                   for k, v in elem.attrib.items())
        else:
            attrs = ""

        if elem.tag in self.self_closing:
            self.add("%s<%s%s />" % (space, elem.tag, attrs))
        else:
            self.add("%s<%s%s>" % (space, elem.tag, attrs))

        for child in elem.xpath("./text()|./node()"):
            if isinstance(child, basestring):
                self.add("%s%s" % (space, child))
            else:
                self.walk(child, indent + 1, stack + [child])
        if elem.tag not in self.self_closing:
            self.add("%s</%s>" % (space, elem.tag))

        if elem.tag == "img":
            self.done()
            # re-open previously closed tags
            for e in stack[:-1]:
                self.add("%s<%s>" % (space, e.tag,))

def node_has_children(elem):
    if elem.xpath("./*"):
        return True
    text_nodes = elem.xpath("./text()")
    text_nodes = [t.strip() for t in text_nodes]
    if any(text_nodes):
        return True
    return False

def postprocess_html_snippet(blob):
    """
    remove empty paragraphs, fix links, ...
    """
    blob = blob.strip()
    if not blob:
        return blob
    try:
        root = html.fromstring(blob)
    except etree.ParserError, e:
        print >> sys.stderr, e
        print >> sys.stderr, repr(blob.encode("utf8"))
        return blob
    for p in root.xpath("//p"):
        if not node_has_children(p):
            p.getparent().remove(p)
    for div in root.xpath("//div"):
        if not node_has_children(div):
            div.getparent().remove(div)
    for a in root.xpath("//a"):
        if not 'href' in a.attrib:
            continue
        orig_href = a.attrib['href']
        if "exit.php" in orig_href or "archive" in orig_href:
            a.attrib['href'] = resolve_url(orig_href)
    return html.tostring(root)

def process_image_html_blob(blob):
    root = html.fromstring(blob)
    img = root.xpath("//img")[0]
    src = img.attrib['src']
    fp = fetch_image(src)
    return fp, {"title": img.attrib.get("title", "") or img.attrib.get("alt", "")}

def image_extension(fp):
    from PIL import Image
    i = Image.open(fp)
    return i.format.lower()

def process_html_blob(blob, cms_content):
    # - h1, h2 need to be mapped to h3, h4 (only one blog post affected, so do manually)
    # - need to handle url mapping
    #   old: https://starship-factory.ch/archives/27-Starship-Factory-heute-n....html
    #   new: /2014/05/01/Starship-Factory-heute-nachmittag-geoeffnet/
    root = html.fromstring(clean_html(blob))

    walker = HTMLWalker()
    walker.walk(root)
    walker.done()

    order = 0
    for snippet_type, snippet in walker.snippets:
        joined = "".join(snippet)
        if not joined.strip():
            continue
        processed = postprocess_html_snippet(joined).encode("utf8")
        if snippet_type == "text":
            ts = TextSnippet()
            ts.content = cms_content
            ts.order = order
            ts.body = processed
            ts.block_slug = "content"
            ts.save()
        if snippet_type == "image":
            fp, img_attrs = process_image_html_blob(processed)
            ext = image_extension(fp)
            isnippet = ImageSnippet()
            isnippet.content = cms_content
            isnippet.order = order
            isnippet.title = unicode(img_attrs['title']) or unicode(fp.name)
            image_filename = "%s.%s" % (my_slugify(isnippet.title), ext)
            isnippet.size = 'content'
            isnippet.image = ImageFile(fp, name=image_filename)
            isnippet.block_slug = "content"
            isnippet.save()
        order += 1

def parse_rss_date(date_as_str):
    dt = datetime.datetime.fromtimestamp(time.mktime(rfc822.parsedate(date_as_str)))
    return make_aware(dt, utc)

@transaction.atomic
def import_from_rss(rss_as_bytes):
    rss = objectify.fromstring(rss_as_bytes)
    current_site = Site.objects.get_current()
    for item in rss.channel.item:
        blob = unicode(item['{%s}encoded' % item.nsmap['content']])
        category, _ = Category.objects.get_or_create(
            slug=my_slugify(unicode(item.category.text)),
            title=unicode(item.category.text)
        )

        entry = Entry()
        cms_content = Content()
        cms_content.save()
        process_html_blob(blob, cms_content)
        entry.content_ptr = cms_content
        entry.title = unicode(item.title.text)
        entry.creation_date = parse_rss_date(unicode(item.pubDate))
        entry.slug = my_slugify(entry.title)
        entry.save()
        entry.sites.add(current_site)
        entry.categories.add(category)
        link = unicode(item.link.text)
        _url_cache[link] = entry.get_absolute_url()
        _url_cache[re.sub("^https", "http", link)] = entry.get_absolute_url()
        add_redirect(link, entry.get_absolute_url())

def main(url_or_fname):
    global _url_cache
    if os.path.exists("url_cache.json"):
	with open("url_cache.json") as fp:
		_url_cache = json.load(fp)

    try:
        if url_or_fname.startswith('http'):
            resp = requests.get(url_or_fname)
            import_from_rss(resp.content)
        else:
            import_from_rss(open(url_or_fname).read())
    finally:
        with open("url_cache.json", "w") as fp:
            json.dump(_url_cache, fp)

if __name__ == "__main__":
    main(sys.argv[1])
<commit_msg>isdir is found in os.path, not os.<commit_after># -*- encoding: utf-8 -*-
"""
Import from s9y RSS export

Warning: this is specifically written for

https://starship-factory.ch/rss.php?version=2.0&all=1

and is pretty much a one-off script. Don't expect it to work on
other RSS feeds, it might not work or break your site/eat your babies/...

"""
import os
import re
import sys
import json
import time
import rfc822
import hashlib
import urlparse
import datetime
from contextlib import contextmanager
import requests
from lxml import objectify, html, etree
from lxml.html.clean import clean_html

from django.db import transaction
from django.utils.html import escape
from django.utils.text import slugify
from django.utils.timezone import make_aware, utc
from django.core.files.images import ImageFile
from django.contrib.sites.models import Site
from django.contrib.redirects.models import Redirect

from simplecms.models import TextSnippet, ImageSnippet, Content
from zinnia.models.entry import Entry
from zinnia.models.category import Category

_url_cache = {}

def my_slugify(value):
    """
    slugify with umlauts
    """
    translation_table = {
        u'ß': 'ss',
        u'ä': 'ae',
        u'ö': 'oe',
        u'ü': 'ue',
        u'é': 'e',
        u'Ä': 'Ae',
        u'Ö': 'Oe',
        u'Ü': 'Ue',
        u'É': 'E',
    }
    for k, v in translation_table.items():
        value = value.replace(k, v)
    return slugify(value)

def follow_redirects(url, verify=True):
    """
    follow any redirects and return the resulting url
    """
    s = requests.Session()
    s.max_redirects = 10  # FIXME: magic number...
    s.verify = verify
    try:
        resp = s.head(url, allow_redirects=True)
    except requests.exceptions.SSLError, e:
        if verify:
            return follow_redirects(url, verify=False)
        else:
            print >> sys.stderr, e, url
            return url
    except requests.exceptions.ConnectionError, e:
        print >> sys.stderr, e, url
        return url
    return resp.url

def transform_image_url(url):
    """
    rewrite the url to the image to one that is closer to the original size
    """
    from django.utils.http import urlencode
    url = url.replace("serendipityThumb.", "")
    if "admin_image_selector.php" in url:
        width_key = 'serendipity[resizeWidth]'
        height_key = 'serendipity[resizeHeight]'
        parsed_url = urlparse.urlparse(url)
        qs = parsed_url.query
        parsed_qs = urlparse.parse_qs(qs)

        if height_key in parsed_qs:
            parsed_qs[height_key] = "10000"
        if width_key in parsed_qs:
            parsed_qs[width_key] = "10000"

        qs = urlencode(parsed_qs, doseq=1)
        url = parsed_url._replace(query=qs).geturl()
    return url

def add_redirect(full_old_url, new_path):
    """
    add a redirect entry. for example
    full_old_url=http://example.org/blog/post1/
    new_path=/2014/05/06/post1/
    """
    current_site = Site.objects.get_current()
    parsed_old = urlparse.urlparse(full_old_url)
    if parsed_old.query:
        # bail out: can only handle paths
        return
    # theres a unique constraint on (site_id, old_path), so get rid of old redirects
    Redirect.objects.filter(site=current_site,
                            old_path=parsed_old.path).delete()
    Redirect.objects.get_or_create(site=current_site,
                                   old_path=parsed_old.path,
                                   new_path=new_path)

def image_cache_path(url):
    key = hashlib.md5(url).hexdigest()
    if not os.path.isdir("image_cache"):
        os.mkdir("image_cache")
    return os.path.join("image_cache", key)

def image_cache_get(url):
    """
    get a file handle for the image at ``url``
    """
    path = image_cache_path(url)
    if os.path.exists(path):
        return open(path)

@contextmanager
def image_cache_put(url):
    path = image_cache_path(url)
    fp = open(path, "wb")
    yield fp
    fp.close()

def fetch_image(url):
    url = transform_image_url(url)
    img = image_cache_get(url)
    if img:
        return img
    resp = requests.get(url, stream=True)
    if resp.status_code == 200:
        with image_cache_put(url) as fp:
            for chunk in resp.iter_content(1024):
                fp.write(chunk)
    return image_cache_get(url)

def resolve_url(url):
    if url in _url_cache:
        return _url_cache[url]
    dest_url = follow_redirects(url)
    _url_cache[url] = dest_url
    return dest_url

class HTMLWalker(object):
    """
    walk an element tree; when encountering an image, save the elements encountered
    from ``self.current_snippet`` into ``self.snippets`` and start a new snippet.
    """
    self_closing = ["img"]

    def __init__(self):
        self.current_snippet = []
        self.snippets = []
        self.current_type = "text"

    def add(self, blob):
        self.current_snippet.append(blob)

    def done(self, new_type="text"):
        self.snippets.append((self.current_type, self.current_snippet))
        self.current_type = new_type
        self.current_snippet = []

    def walk(self, elem, indent=0, stack=None):
        stack = stack or []
        space = indent * 4 * " "
        space = ""

        if elem.tag == "img":
            if "src" not in elem.attrib:
                return
            # close any open tags
            for e in list(reversed(stack))[1:]:
                self.add("%s</%s>" % (space, e.tag,))
            self.done(new_type="image")

        # special case: thumbnail linked to the full image
        if elem.tag == "a" and elem.attrib.get("class", "") == "serendipity_image_link":
            self.done(new_type="image")
            title = elem.find("img").attrib.get("title", "")
            self.add("" %
                     (elem.attrib.get("href"),
                      title, title))
            self.done()
            return

        if elem.tag == "div" and len(elem.xpath("//*")) == 0:
            return

        if elem.attrib:
            attrs = " " + " ".join('%s="%s"' % (escape(k), escape(v))
                                   for k, v in elem.attrib.items())
        else:
            attrs = ""

        if elem.tag in self.self_closing:
            self.add("%s<%s%s />" % (space, elem.tag, attrs))
        else:
            self.add("%s<%s%s>" % (space, elem.tag, attrs))

        for child in elem.xpath("./text()|./node()"):
            if isinstance(child, basestring):
                self.add("%s%s" % (space, child))
            else:
                self.walk(child, indent + 1, stack + [child])
        if elem.tag not in self.self_closing:
            self.add("%s</%s>" % (space, elem.tag))

        if elem.tag == "img":
            self.done()
            # re-open previously closed tags
            for e in stack[:-1]:
                self.add("%s<%s>" % (space, e.tag,))

def node_has_children(elem):
    if elem.xpath("./*"):
        return True
    text_nodes = elem.xpath("./text()")
    text_nodes = [t.strip() for t in text_nodes]
    if any(text_nodes):
        return True
    return False

def postprocess_html_snippet(blob):
    """
    remove empty paragraphs, fix links, ...
    """
    blob = blob.strip()
    if not blob:
        return blob
    try:
        root = html.fromstring(blob)
    except etree.ParserError, e:
        print >> sys.stderr, e
        print >> sys.stderr, repr(blob.encode("utf8"))
        return blob
    for p in root.xpath("//p"):
        if not node_has_children(p):
            p.getparent().remove(p)
    for div in root.xpath("//div"):
        if not node_has_children(div):
            div.getparent().remove(div)
    for a in root.xpath("//a"):
        if not 'href' in a.attrib:
            continue
        orig_href = a.attrib['href']
        if "exit.php" in orig_href or "archive" in orig_href:
            a.attrib['href'] = resolve_url(orig_href)
    return html.tostring(root)

def process_image_html_blob(blob):
    root = html.fromstring(blob)
    img = root.xpath("//img")[0]
    src = img.attrib['src']
    fp = fetch_image(src)
    return fp, {"title": img.attrib.get("title", "") or img.attrib.get("alt", "")}

def image_extension(fp):
    from PIL import Image
    i = Image.open(fp)
    return i.format.lower()

def process_html_blob(blob, cms_content):
    # - h1, h2 need to be mapped to h3, h4 (only one blog post affected, so do manually)
    # - need to handle url mapping
    #   old: https://starship-factory.ch/archives/27-Starship-Factory-heute-n....html
    #   new: /2014/05/01/Starship-Factory-heute-nachmittag-geoeffnet/
    root = html.fromstring(clean_html(blob))

    walker = HTMLWalker()
    walker.walk(root)
    walker.done()

    order = 0
    for snippet_type, snippet in walker.snippets:
        joined = "".join(snippet)
        if not joined.strip():
            continue
        processed = postprocess_html_snippet(joined).encode("utf8")
        if snippet_type == "text":
            ts = TextSnippet()
            ts.content = cms_content
            ts.order = order
            ts.body = processed
            ts.block_slug = "content"
            ts.save()
        if snippet_type == "image":
            fp, img_attrs = process_image_html_blob(processed)
            ext = image_extension(fp)
            isnippet = ImageSnippet()
            isnippet.content = cms_content
            isnippet.order = order
            isnippet.title = unicode(img_attrs['title']) or unicode(fp.name)
            image_filename = "%s.%s" % (my_slugify(isnippet.title), ext)
            isnippet.size = 'content'
            isnippet.image = ImageFile(fp, name=image_filename)
            isnippet.block_slug = "content"
            isnippet.save()
        order += 1

def parse_rss_date(date_as_str):
    dt = datetime.datetime.fromtimestamp(time.mktime(rfc822.parsedate(date_as_str)))
    return make_aware(dt, utc)

@transaction.atomic
def import_from_rss(rss_as_bytes):
    rss = objectify.fromstring(rss_as_bytes)
    current_site = Site.objects.get_current()
    for item in rss.channel.item:
        blob = unicode(item['{%s}encoded' % item.nsmap['content']])
        category, _ = Category.objects.get_or_create(
            slug=my_slugify(unicode(item.category.text)),
            title=unicode(item.category.text)
        )

        entry = Entry()
        cms_content = Content()
        cms_content.save()
        process_html_blob(blob, cms_content)
        entry.content_ptr = cms_content
        entry.title = unicode(item.title.text)
        entry.creation_date = parse_rss_date(unicode(item.pubDate))
        entry.slug = my_slugify(entry.title)
        entry.save()
        entry.sites.add(current_site)
        entry.categories.add(category)
        link = unicode(item.link.text)
        _url_cache[link] = entry.get_absolute_url()
        _url_cache[re.sub("^https", "http", link)] = entry.get_absolute_url()
        add_redirect(link, entry.get_absolute_url())

def main(url_or_fname):
    global _url_cache
    if os.path.exists("url_cache.json"):
	with open("url_cache.json") as fp:
		_url_cache = json.load(fp)

    try:
        if url_or_fname.startswith('http'):
            resp = requests.get(url_or_fname)
            import_from_rss(resp.content)
        else:
            import_from_rss(open(url_or_fname).read())
    finally:
        with open("url_cache.json", "w") as fp:
            json.dump(_url_cache, fp)

if __name__ == "__main__":
    main(sys.argv[1])
<|endoftext|>