<commit_before>from data_collection.data_types import Address, AddressSet
from django.test import TestCase

class MockLogger:
    def log_message(self, level, message, variable=None, pretty=False):
        pass

class AddressSetTest(TestCase):

    def test_add_with_duplicates(self):
        in_list = [
            {'address': 'foo', 'slug': 'foo', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
            {'address': 'bar', 'slug': 'bar', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
            {'address': 'foo', 'slug': 'foo', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
        ]

        expected = set([
            Address(
                address='foo', slug='foo', postcode='', council='', polling_station_id='', uprn='', location=None),
            Address(
                address='bar', slug='bar', postcode='', council='', polling_station_id='', uprn='', location=None),
        ])

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)

        self.assertEqual(expected, address_list.elements)

    def test_remove_ambiguous_addresses_exactmatch(self):
        in_list = [
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'haringey-park-london-n89jg-aa',
                'uprn': '',
            },
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'haringey-park-london-n89jg-ab',
                'uprn': '',
            },
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AC',
                'council': '',
                'slug': 'haringey-park-london-n89jg-ac',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_fuzzymatch(self):
        """
        The addresses:
        - 5-6 Mickleton Dr, Southport, PR82QX
        - 5/6, Mickleton Dr.  Southport, PR82QX
        - 5-6 mickleton dr southport, PR82QX
        should all be considered the same address.

        The addresses:
        - 56 Mickleton Dr, Southport, PR82QX
        - 5-6 Mickleton Dr, Southport, BT281QZ
        should be identified as different.
        """
        in_list = [
            {
                'address': '5-6 Mickleton Dr, Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A01',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '5/6, Mickleton Dr.  Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A02',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '5-6 mickleton dr southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A03',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '56 Mickleton Dr, Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A04',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
            {
                'address': '5-6 Mickleton Dr, Southport',
                'postcode': 'BT281QZ',
                'polling_station_id': 'A05',
                'council': '',
                'slug': 'e',
                'uprn': '',
            },
            {
                'address': '56 Mickleton Dr, Southport',
                'postcode': 'BT281QZ',
                'polling_station_id': 'A04',
                'council': '',
                'slug': 'f',
                'uprn': '',
            },
        ]

        expected = set([
            Address(
                address='5-6 Mickleton Dr, Southport',
                postcode='BT281QZ',
                polling_station_id='A05',
                council='',
                slug='e',
                uprn='',
                location=None),
            Address(
                address='56 Mickleton Dr, Southport',
                postcode='BT281QZ',
                polling_station_id='A04',
                council='',
                slug='f',
                uprn='',
                location=None),
        ])

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(expected, result)

    def test_remove_ambiguous_addresses_some_stations_match(self):
        # if one polling station doesn't match, we should remove all of them
        in_list = [
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_whole_postcode(self):
        # if we've got one ambiguous address,
        # we should remove all addresse with the same postcode
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_no_issues(self):
        # if there are no ambiguous addresses,
        # we shouldn't do anything

        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        expected = set(address_list.elements)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(expected, result)
<commit_msg>add unit tests for new AddressSet methods<commit_after>from data_collection.data_types import Address, AddressSet
from django.test import TestCase

class MockLogger:
    def log_message(self, level, message, variable=None, pretty=False):
        pass

class AddressSetTest(TestCase):

    def test_add_with_duplicates(self):
        in_list = [
            {'address': 'foo', 'slug': 'foo', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
            {'address': 'bar', 'slug': 'bar', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
            {'address': 'foo', 'slug': 'foo', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
        ]

        expected = set([
            Address(
                address='foo', slug='foo', postcode='', council='', polling_station_id='', uprn='', location=None),
            Address(
                address='bar', slug='bar', postcode='', council='', polling_station_id='', uprn='', location=None),
        ])

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)

        self.assertEqual(expected, address_list.elements)

    def test_remove_ambiguous_addresses_exactmatch(self):
        in_list = [
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'haringey-park-london-n89jg-aa',
                'uprn': '',
            },
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'haringey-park-london-n89jg-ab',
                'uprn': '',
            },
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AC',
                'council': '',
                'slug': 'haringey-park-london-n89jg-ac',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_fuzzymatch(self):
        """
        The addresses:
        - 5-6 Mickleton Dr, Southport, PR82QX
        - 5/6, Mickleton Dr.  Southport, PR82QX
        - 5-6 mickleton dr southport, PR82QX
        should all be considered the same address.

        The addresses:
        - 56 Mickleton Dr, Southport, PR82QX
        - 5-6 Mickleton Dr, Southport, BT281QZ
        should be identified as different.
        """
        in_list = [
            {
                'address': '5-6 Mickleton Dr, Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A01',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '5/6, Mickleton Dr.  Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A02',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '5-6 mickleton dr southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A03',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '56 Mickleton Dr, Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A04',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
            {
                'address': '5-6 Mickleton Dr, Southport',
                'postcode': 'BT281QZ',
                'polling_station_id': 'A05',
                'council': '',
                'slug': 'e',
                'uprn': '',
            },
            {
                'address': '56 Mickleton Dr, Southport',
                'postcode': 'BT281QZ',
                'polling_station_id': 'A04',
                'council': '',
                'slug': 'f',
                'uprn': '',
            },
        ]

        expected = set([
            Address(
                address='5-6 Mickleton Dr, Southport',
                postcode='BT281QZ',
                polling_station_id='A05',
                council='',
                slug='e',
                uprn='',
                location=None),
            Address(
                address='56 Mickleton Dr, Southport',
                postcode='BT281QZ',
                polling_station_id='A04',
                council='',
                slug='f',
                uprn='',
                location=None),
        ])

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(expected, result)

    def test_remove_ambiguous_addresses_some_stations_match(self):
        # if one polling station doesn't match, we should remove all of them
        in_list = [
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_whole_postcode(self):
        # if we've got one ambiguous address,
        # we should remove all addresse with the same postcode
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_no_issues(self):
        # if there are no ambiguous addresses,
        # we shouldn't do anything

        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        expected = set(address_list.elements)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(expected, result)

    def test_attach_doorstep_grid_refs(self):
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '00001',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '00004',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)

        addressbase = {
            # 00001 is in here but 00004 isn't
            '00001': {'location': "SRID=4326;POINT(-0.9288492 53.3119342)"}
        }

        address_set.elements = address_set.attach_doorstep_gridrefs(addressbase)

        self.assertEqual(2, len(address_set.elements))
        for el in address_set.elements:
            if el.uprn == '00001':
                self.assertEqual("SRID=4326;POINT(-0.9288492 53.3119342)", el.location)
            if el.uprn == '00004':
                self.assertEqual(None, el.location)

    def test_remove_invalid_uprns(self):
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '00001',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '00002',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '00003',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '00004',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)

        addressbase = {
            '00001': {
                'postcode': 'L252NW',
                'address': '1 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288492 53.3119342)",
            },
            '00002': {
                'postcode': 'L252NW',
                'address': '2 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288480 53.3119345)",
            },
            '00003': {
                'postcode': 'L252XX',  # this postcode doesn't match with the input record
                'address': '2 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288400 53.3119332)",
            }
            # 00004 is not in here
        }

        address_set.elements = address_set.remove_invalid_uprns(addressbase)

        # 00003 and 00004 should still be in the set
        self.assertEqual(4, len(address_set.elements))
        # but those records should now have a blank uprn
        for el in address_set.elements:
            assert el.uprn in ['00001', '00002', '']
<|endoftext|>


<commit_before>from data_collection.data_types import Address, AddressSet
from django.test import TestCase

class MockLogger:
    def log_message(self, level, message, variable=None, pretty=False):
        pass

class AddressSetTest(TestCase):

    def test_add_with_duplicates(self):
        in_list = [
            {'address': 'foo', 'slug': 'foo', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
            {'address': 'bar', 'slug': 'bar', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
            {'address': 'foo', 'slug': 'foo', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
        ]

        expected = set([
            Address(
                address='foo', slug='foo', postcode='', council='', polling_station_id='', uprn='', location=None),
            Address(
                address='bar', slug='bar', postcode='', council='', polling_station_id='', uprn='', location=None),
        ])

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)

        self.assertEqual(expected, address_list.elements)

    def test_remove_ambiguous_addresses_exactmatch(self):
        in_list = [
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'haringey-park-london-n89jg-aa',
                'uprn': '',
            },
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'haringey-park-london-n89jg-ab',
                'uprn': '',
            },
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AC',
                'council': '',
                'slug': 'haringey-park-london-n89jg-ac',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_fuzzymatch(self):
        """
        The addresses:
        - 5-6 Mickleton Dr, Southport, PR82QX
        - 5/6, Mickleton Dr.  Southport, PR82QX
        - 5-6 mickleton dr southport, PR82QX
        should all be considered the same address.

        The addresses:
        - 56 Mickleton Dr, Southport, PR82QX
        - 5-6 Mickleton Dr, Southport, BT281QZ
        should be identified as different.
        """
        in_list = [
            {
                'address': '5-6 Mickleton Dr, Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A01',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '5/6, Mickleton Dr.  Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A02',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '5-6 mickleton dr southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A03',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '56 Mickleton Dr, Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A04',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
            {
                'address': '5-6 Mickleton Dr, Southport',
                'postcode': 'BT281QZ',
                'polling_station_id': 'A05',
                'council': '',
                'slug': 'e',
                'uprn': '',
            },
            {
                'address': '56 Mickleton Dr, Southport',
                'postcode': 'BT281QZ',
                'polling_station_id': 'A04',
                'council': '',
                'slug': 'f',
                'uprn': '',
            },
        ]

        expected = set([
            Address(
                address='5-6 Mickleton Dr, Southport',
                postcode='BT281QZ',
                polling_station_id='A05',
                council='',
                slug='e',
                uprn='',
                location=None),
            Address(
                address='56 Mickleton Dr, Southport',
                postcode='BT281QZ',
                polling_station_id='A04',
                council='',
                slug='f',
                uprn='',
                location=None),
        ])

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(expected, result)

    def test_remove_ambiguous_addresses_some_stations_match(self):
        # if one polling station doesn't match, we should remove all of them
        in_list = [
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_whole_postcode(self):
        # if we've got one ambiguous address,
        # we should remove all addresse with the same postcode
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_no_issues(self):
        # if there are no ambiguous addresses,
        # we shouldn't do anything

        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_list = AddressSet(MockLogger())
        for el in in_list:
            address_list.add(el)
        expected = set(address_list.elements)
        result = address_list.remove_ambiguous_addresses()

        self.assertEqual(expected, result)

    def test_attach_doorstep_grid_refs(self):
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '00001',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '00004',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)

        addressbase = {
            # 00001 is in here but 00004 isn't
            '00001': {'location': "SRID=4326;POINT(-0.9288492 53.3119342)"}
        }

        address_set.elements = address_set.attach_doorstep_gridrefs(addressbase)

        self.assertEqual(2, len(address_set.elements))
        for el in address_set.elements:
            if el.uprn == '00001':
                self.assertEqual("SRID=4326;POINT(-0.9288492 53.3119342)", el.location)
            if el.uprn == '00004':
                self.assertEqual(None, el.location)

    def test_remove_invalid_uprns(self):
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '00001',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '00002',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '00003',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '00004',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)

        addressbase = {
            '00001': {
                'postcode': 'L252NW',
                'address': '1 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288492 53.3119342)",
            },
            '00002': {
                'postcode': 'L252NW',
                'address': '2 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288480 53.3119345)",
            },
            '00003': {
                'postcode': 'L252XX',  # this postcode doesn't match with the input record
                'address': '2 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288400 53.3119332)",
            }
            # 00004 is not in here
        }

        address_set.elements = address_set.remove_invalid_uprns(addressbase)

        # 00003 and 00004 should still be in the set
        self.assertEqual(4, len(address_set.elements))
        # but those records should now have a blank uprn
        for el in address_set.elements:
            assert el.uprn in ['00001', '00002', '']
<commit_msg>fix inaccurate var names<commit_after>from data_collection.data_types import Address, AddressSet
from django.test import TestCase

class MockLogger:
    def log_message(self, level, message, variable=None, pretty=False):
        pass

class AddressSetTest(TestCase):

    def test_add_with_duplicates(self):
        in_list = [
            {'address': 'foo', 'slug': 'foo', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
            {'address': 'bar', 'slug': 'bar', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
            {'address': 'foo', 'slug': 'foo', 'postcode': '', 'council': '', 'polling_station_id': '', 'uprn': '',},
        ]

        expected = set([
            Address(
                address='foo', slug='foo', postcode='', council='', polling_station_id='', uprn='', location=None),
            Address(
                address='bar', slug='bar', postcode='', council='', polling_station_id='', uprn='', location=None),
        ])

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)

        self.assertEqual(expected, address_set.elements)

    def test_remove_ambiguous_addresses_exactmatch(self):
        in_list = [
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'haringey-park-london-n89jg-aa',
                'uprn': '',
            },
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'haringey-park-london-n89jg-ab',
                'uprn': '',
            },
            {
                'address': 'Haringey Park, London',
                'postcode': 'N89JG',
                'polling_station_id': 'AC',
                'council': '',
                'slug': 'haringey-park-london-n89jg-ac',
                'uprn': '',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)
        result = address_set.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_fuzzymatch(self):
        """
        The addresses:
        - 5-6 Mickleton Dr, Southport, PR82QX
        - 5/6, Mickleton Dr.  Southport, PR82QX
        - 5-6 mickleton dr southport, PR82QX
        should all be considered the same address.

        The addresses:
        - 56 Mickleton Dr, Southport, PR82QX
        - 5-6 Mickleton Dr, Southport, BT281QZ
        should be identified as different.
        """
        in_list = [
            {
                'address': '5-6 Mickleton Dr, Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A01',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '5/6, Mickleton Dr.  Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A02',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '5-6 mickleton dr southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A03',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '56 Mickleton Dr, Southport',
                'postcode': 'PR82QX',
                'polling_station_id': 'A04',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
            {
                'address': '5-6 Mickleton Dr, Southport',
                'postcode': 'BT281QZ',
                'polling_station_id': 'A05',
                'council': '',
                'slug': 'e',
                'uprn': '',
            },
            {
                'address': '56 Mickleton Dr, Southport',
                'postcode': 'BT281QZ',
                'polling_station_id': 'A04',
                'council': '',
                'slug': 'f',
                'uprn': '',
            },
        ]

        expected = set([
            Address(
                address='5-6 Mickleton Dr, Southport',
                postcode='BT281QZ',
                polling_station_id='A05',
                council='',
                slug='e',
                uprn='',
                location=None),
            Address(
                address='56 Mickleton Dr, Southport',
                postcode='BT281QZ',
                polling_station_id='A04',
                council='',
                slug='f',
                uprn='',
                location=None),
        ])

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)
        result = address_set.remove_ambiguous_addresses()

        self.assertEqual(expected, result)

    def test_remove_ambiguous_addresses_some_stations_match(self):
        # if one polling station doesn't match, we should remove all of them
        in_list = [
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': 'Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)
        result = address_set.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_whole_postcode(self):
        # if we've got one ambiguous address,
        # we should remove all addresse with the same postcode
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)
        result = address_set.remove_ambiguous_addresses()

        self.assertEqual(set(), result)

    def test_remove_ambiguous_addresses_no_issues(self):
        # if there are no ambiguous addresses,
        # we shouldn't do anything

        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)
        expected = set(address_set.elements)
        result = address_set.remove_ambiguous_addresses()

        self.assertEqual(expected, result)

    def test_attach_doorstep_grid_refs(self):
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '00001',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '00004',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)

        addressbase = {
            # 00001 is in here but 00004 isn't
            '00001': {'location': "SRID=4326;POINT(-0.9288492 53.3119342)"}
        }

        address_set.elements = address_set.attach_doorstep_gridrefs(addressbase)

        self.assertEqual(2, len(address_set.elements))
        for el in address_set.elements:
            if el.uprn == '00001':
                self.assertEqual("SRID=4326;POINT(-0.9288492 53.3119342)", el.location)
            if el.uprn == '00004':
                self.assertEqual(None, el.location)

    def test_remove_invalid_uprns(self):
        in_list = [
            {
                'address': '1 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'a',
                'uprn': '00001',
            },
            {
                'address': '2 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'b',
                'uprn': '00002',
            },
            {
                'address': '3 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AB',
                'council': '',
                'slug': 'c',
                'uprn': '00003',
            },
            {
                'address': '4 Abbeyvale Dr, Liverpool',
                'postcode': 'L252NW',
                'polling_station_id': 'AA',
                'council': '',
                'slug': 'd',
                'uprn': '00004',
            },
        ]

        address_set = AddressSet(MockLogger())
        for el in in_list:
            address_set.add(el)

        addressbase = {
            '00001': {
                'postcode': 'L252NW',
                'address': '1 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288492 53.3119342)",
            },
            '00002': {
                'postcode': 'L252NW',
                'address': '2 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288480 53.3119345)",
            },
            '00003': {
                'postcode': 'L252XX',  # this postcode doesn't match with the input record
                'address': '2 Abbeyvale Dr, Liverpool',
                'location': "SRID=4326;POINT(-0.9288400 53.3119332)",
            }
            # 00004 is not in here
        }

        address_set.elements = address_set.remove_invalid_uprns(addressbase)

        # 00003 and 00004 should still be in the set
        self.assertEqual(4, len(address_set.elements))
        # but those records should now have a blank uprn
        for el in address_set.elements:
            assert el.uprn in ['00001', '00002', '']
<|endoftext|>


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

import argparse
from subprocess import call

def run(command):
    print command
    call(command.split())

def docker_run(args):
    if args == None:
        return
    if args.container == 'firefox':
        run('docker run --name firefox -p {vnc}:5901 -p {rdp}:3389 -e "GEOMETRY={g}" rickdesantis/firefox'.format(vnc=args.vnc, rdp=args.rdp, g=args.g))
    elif args.container == 'centos':
        run('docker run --name centos rickdesantis/centos')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Runs the containers uploaded by the docker hub user rickdesantis.')
    parser.add_argument("container", help='the name of the container that will be run', choices=['firefox', 'centos'])
    parser.add_argument("-g", help='the screen resolution that will be used', default='1440x900')
    parser.add_argument("-vnc", help='the port used by VNC', default='5901')
    parser.add_argument("-rdp", help='the port used by RDP', default='3389')
    args = parser.parse_args()
    docker_run(args)
<commit_msg>Updated the script<commit_after>#!/usr/bin/env python

import argparse
import subprocess

def run(command, background=False):
    print command
    if background:
        subprocess.Popen(command.split())
    else:
        subprocess.call(command.split())

def docker_run(args):
    if args == None:
        return
    if args.container == 'firefox':
        run('docker stop firefox')
        run('docker rm firefox')
        if not args.stop:
            run('docker run --name firefox -p {vnc}:5901 -p {rdp}:3389 -e "GEOMETRY={g}" rickdesantis/firefox'.format(vnc=args.vnc, rdp=args.rdp, g=args.g), True)
            if not args.startonly:
                run('open vnc://127.0.0.1:5901')
    elif args.container == 'centos':
        run('docker stop centos')
        run('docker rm centos')
        if not args.stop:
            run('docker run --name centos rickdesantis/centos', True)
            if not args.startonly:
                run('docker exec -it centos bash')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Runs the containers uploaded by the docker hub user rickdesantis.')
    parser.add_argument("container", help='the name of the container that will be run', choices=['firefox', 'centos'])
    parser.add_argument("-g", help='the screen resolution that will be used', default='1440x900')
    parser.add_argument("-vnc", help='the port used by VNC', default='5901')
    parser.add_argument("-rdp", help='the port used by RDP', default='3389')
    parser.add_argument("-stop", help='stops the container', action='store_true')
    parser.add_argument("-startonly", help='starts the container but don\'t connect to it', action='store_true')
    args = parser.parse_args()
    docker_run(args)
<|endoftext|>


<commit_before>import os
import shutil
import subprocess
import sys
import random
import threading
import json
import time
import traceback
import signal

USER_HOME_DIR = '~'
HOME_DIR = '~/Docker'
APP_DIR = '~/.local/share/applications'
PULSE_COOKIE_PATH = '~/.config/pulse/cookie'
TMP_DIR = '/tmp'
ROOT_DIR = os.path.dirname(os.path.realpath(__file__))
SCRIPT_PATH = os.path.join(ROOT_DIR, __file__)
PULSE_SERVER = 'unix:/var/run/pulse/native'

# Init
DESKTOP_ENTRY = '''
[Desktop Entry]
Version=1.0
Type=Application
Terminal=false
Name=Docker - %s
Comment=Docker - %s
Exec=%s
Icon=%s
Categories=Other;
'''
USER_HOME_DIR = os.path.expanduser(USER_HOME_DIR)
HOME_DIR = os.path.expanduser(HOME_DIR)
DESKTOP_DIR = os.path.expanduser(APP_DIR)
PULSE_COOKIE_PATH = os.path.expanduser(PULSE_COOKIE_PATH)
TMP_DIR = os.path.expanduser(TMP_DIR)
CONF_DIR = os.path.join(HOME_DIR, '.config')
BASE_CONF_PATH = os.path.join(CONF_DIR, 'base.json')
interrupt = False

def mkdirs(path):
    if not os.path.exists(path):
        os.makedirs(path)

mkdirs(CONF_DIR)

if not os.path.exists(BASE_CONF_PATH):
    shutil.copyfile(
        os.path.join(ROOT_DIR, 'apps', 'base', 'settings.json'),
        BASE_CONF_PATH,
    )

with open(BASE_CONF_PATH, 'r') as conf_file:
    conf_data = json.loads(conf_file.read())

SUDO_DOCKER = conf_data.get('sudo_docker', True)
INCREASE_SHM = conf_data.get('increase_shm', True)
SHARE_CLIPBOARD = conf_data.get('share_clipboard', True)
SHARE_FONTS = conf_data.get('share_fonts', True)
SHARE_THEMES = conf_data.get('share_themes', False)
SHARE_ICONS = conf_data.get('share_icons', False)
SHARE_USER_FONTS = conf_data.get('share_user_fonfs', True)
SHARE_USER_THEMES = conf_data.get('share_user_themes', True)
DPI = conf_data.get('dpi')
DEBUG = False

GPU = conf_data.get('gpu', 'auto')
if GPU == 'auto':
    try:
        subprocess.check_output(['which', 'nvidia-settings'],
            stderr=subprocess.PIPE)
        GPU = 'nvidia'
    except:
        GPU = 'intel'

DEFAULT_WIN_SIZE = conf_data.get('default_win_size', '1024x768')
DEFAULT_VOLUMES = conf_data.get('default_volumes', [])

def kill_process(process):
    terminated = False

    for _ in xrange(200):
        try:
            process.send_signal(signal.SIGINT)
        except OSError as error:
            if error.errno != 3:
                raise
        for _ in xrange(4):
            if process.poll() is not None:
                terminated = True
                break
            time.sleep(0.0025)
        if terminated:
            break

    if not terminated:
        for _ in xrange(10):
            if process.poll() is not None:
                break
            try:
                process.send_signal(signal.SIGKILL)
            except OSError as error:
                if error.errno != 3:
                    raise
            time.sleep(0.01)

def image_exists(image):
    image_id = subprocess.check_output([
        'docker',
        'images',
        '-q',
        image,
    ]).strip()

    return bool(image_id)

def pull():
    subprocess.check_call((['sudo'] if SUDO_DOCKER else []) + [
        'docker',
        'pull',
        'pritunl/archlinux',
    ])

def exists_pull():
    if not image_exists('pritunl/archlinux'):
        pull()

def build(app):
    app = app.split('#')[0]

    if app == 'base-intel' or app == 'base-nvidia':
        image_name = 'base-xorg'
    else:
        image_name = app

    app_dir = os.path.join(ROOT_DIR, 'apps', app)

    subprocess.check_call((['sudo'] if SUDO_DOCKER else []) + [
        'docker',
        'build',
        '--rm',
        '-t', 'dockercask/' + image_name,
        '.',
    ], cwd=app_dir)

    if app == 'base':
        build('base-' + GPU)

def exists_build(app):
    app = app.split('#')[0]

    if not image_exists('dockercask/base'):
        build('base')

    if not image_exists('dockercask/' + app):
        build(app)

def build_all():
    build('base')
    for app in os.listdir(HOME_DIR):
        if not os.path.isdir(os.path.join(HOME_DIR, app)) or \
                app.startswith('.'):
            continue
        build(app)

def add(app):
    app_dir = os.path.join(HOME_DIR, app)
    icon_path = os.path.join(ROOT_DIR, 'apps', app.split('#')[0], 'icon.png')
    app_conf_path = os.path.join(CONF_DIR, app + '.json')
    desktop_entry_path = os.path.join(DESKTOP_DIR,
        'docker-%s.desktop' % app.replace('#', '-'))

    mkdirs(app_dir)

    if not os.path.exists(app_conf_path):
        shutil.copyfile(
            os.path.join(ROOT_DIR, 'apps', app.split('#')[0], 'settings.json'),
            app_conf_path,
        )

    if DEBUG:
        cmd = 'xfce4-terminal --command="python2 %s %s --debug"' % (
            SCRIPT_PATH, app)
    else:
        cmd = 'python2 %s %s' % (SCRIPT_PATH, app)

    formated_app_name = app.replace('#', ' ').replace('-', ' ').split()
    formated_app_name = ' '.join([x.capitalize() for x in formated_app_name])

    with open(desktop_entry_path, 'w') as desktop_file:
        desktop_file.write(DESKTOP_ENTRY % (
            formated_app_name,
            formated_app_name,
            cmd,
            icon_path,
        ))

def remove(app):
    app_dir = os.path.join(HOME_DIR, app)
    app_conf_path = os.path.join(CONF_DIR, app + '.json')
    desktop_entry_path = os.path.join(DESKTOP_DIR,
        'docker-%s.desktop' % app.replace('#', '-'))

    for path in (app_dir, app_conf_path, desktop_entry_path):
        subprocess.check_call([
            'rm',
            '-rf',
            path,
        ])

def run(app):
    try:
        subprocess.check_call([
            'wmctrl',
            '-F',
            '-R',
            app,
        ])
        return
    except subprocess.CalledProcessError:
        pass

    app_dir = os.path.join(HOME_DIR, app)
    app_conf_path = os.path.join(CONF_DIR, app + '.json')
    fonts_dir = os.path.join(USER_HOME_DIR, '.fonts')
    themes_dir = os.path.join(USER_HOME_DIR, '.themes')
    cmd = []
    docker_args = []
    volume_args = []

    if not os.path.exists(app_dir):
        print 'App must be added before running'
        exit(1)

    with open(app_conf_path, 'r') as app_conf_file:
        app_conf_data = json.loads(app_conf_file.read())

    host_x11 = app_conf_data.get('host_x11')
    privileged = app_conf_data.get('privileged')
    increase_shm = app_conf_data.get('increase_shm', INCREASE_SHM)
    dpi = app_conf_data.get('dpi', DPI)

    if DEBUG:
        docker_args.append('-it')
        cmd.append('/bin/bash')

    if privileged:
        docker_args += ['--privileged']

    if increase_shm:
        docker_args += ['--shm-size', '1g']

    for src, dest in DEFAULT_VOLUMES:
        volume_args += [
            '-v', '%s:%s' % (os.path.expanduser(src), dest),
        ]

    if SHARE_FONTS:
        volume_args += [
            '-v', '%s:%s' % ('/usr/share/fonts', '/usr/share/fonts:ro'),
        ]

    if SHARE_ICONS:
        volume_args += [
            '-v', '%s:%s' % ('/usr/share/icons', '/usr/share/icons:ro'),
        ]

    if SHARE_THEMES:
        volume_args += [
            '-v', '%s:%s' % ('/usr/share/themes', '/usr/share/themes:ro'),
        ]

    if SHARE_USER_FONTS:
        volume_args += [
            '-v', '%s:%s' % (fonts_dir, '/home/docker/.fonts:ro'),
        ]

    if SHARE_USER_THEMES:
        volume_args += [
            '-v', '%s:%s' % (themes_dir, '/home/docker/.themes:ro'),
        ]

    mkdirs(fonts_dir)
    mkdirs(themes_dir)

    if host_x11:
        x_cookie = subprocess.check_output(['xauth', 'list', ':0']).split()[-1]
        x_num = '0'
    else:
        x_cookie = subprocess.check_output(['mcookie'])
        x_num = str(random.randint(1000, 32000))
        x_auth_path = os.path.join(TMP_DIR, '.X11-docker-' + x_num)

        with open(x_auth_path, 'w') as _:
            pass

        subprocess.check_call([
            'xauth',
            '-f', x_auth_path,
            'add',
            ':' + x_num,
            'MIT-MAGIC-COOKIE-1',
            x_cookie,
        ])

        subprocess.check_call([
            'xauth',
            'add',
            ':' + x_num,
            'MIT-MAGIC-COOKIE-1',
            x_cookie,
        ])

    x_screen_path = os.path.join(TMP_DIR, '.X11-unix', 'X' + x_num)

    x_proc = None
    docker_id = None

    def clean_up():
        global interrupt
        interrupt = True

        if docker_id:
            try:
                subprocess.check_output([
                    'docker',
                    'rm',
                    '-f',
                    docker_id,
                ], stderr=subprocess.PIPE)
            except:
                pass
        if x_proc:
            kill_process(x_proc)
        if not host_x11:
            try:
                subprocess.check_output([
                    'xauth',
                    'remove',
                    ':' + x_num,
                ], stderr=subprocess.PIPE)
            except:
                pass
            try:
                os.remove(x_auth_path)
            except:
                pass
            try:
                os.remove(x_screen_path)
            except:
                pass

    if not host_x11:
        args = [
            'Xephyr',
            '-auth', x_auth_path,
            '-screen', DEFAULT_WIN_SIZE,
            '-title', app,
            '-br',
            '-resizeable',
            '-no-host-grab',
            '-nolisten', 'tcp',
        ]

        if dpi:
            args += ['-dpi', dpi]

        x_proc = subprocess.Popen(args + [':' + x_num])

        def thread_func():
            try:
                x_proc.wait()
            finally:
                clean_up()

        thread = threading.Thread(target=thread_func)
        thread.start()

    args = (['sudo'] if SUDO_DOCKER else []) + [
        'docker',
        'run',
        '--rm' if DEBUG else '--detach',
    ] + docker_args + [
        '-v', '%s:%s' % (x_screen_path, x_screen_path),
        '-v', '%s:%s' % (PULSE_COOKIE_PATH, '/tmp/.pulse-cookie'),
        '-v', '%s:%s' % (BASE_CONF_PATH, '/base.json:ro'),
        '-v', '%s:%s' % (app_conf_path, '/app.json:ro'),
        '-v', '%s:%s' % (app_dir, '/home/docker/Docker'),
        '-v', '/var/run/user/%s/pulse/native:/var/run/pulse/native' % (
            os.getuid()),
    ] + volume_args + [
        '-u', 'docker',
        '-e', 'HOME=/home/docker',
        '-e', 'DISPLAY=:' + x_num,
        '-e', 'XAUTHORITY=/tmp/.Xauth',
        '-e', 'XCOOKIE=' + x_cookie,
        '-e', 'PULSE_SERVER=' + PULSE_SERVER,
        'dockercask/' + app.split('#')[0],
    ] + cmd

    print ' '.join(args)

    if not host_x11 and SHARE_CLIPBOARD:
        thread = threading.Thread(target=share_clipboard, args=(x_num,))
        thread.daemon = True
        thread.start()

    if DEBUG:
        try:
            subprocess.check_call(args)
        finally:
            clean_up()
    else:
        docker_id = subprocess.check_output(args).strip()

        try:
            subprocess.check_call(['docker', 'wait', docker_id])
        finally:
            clean_up()

def set_clipboard(num, val):
    proc = subprocess.Popen(
        ['xsel', '--display', ':' + num, '-b', '-i', '-t', '500'],
        stdin=subprocess.PIPE,
    )
    proc.communicate(input=val)
    proc.wait()

def get_clipboard(num):
    return subprocess.check_output(
        ['xsel', '--display', ':' + num, '-b', '-o', '-t', '500'],
    )

def share_clipboard(app_num):
    time.sleep(1)

    val = get_clipboard('0')
    set_clipboard(app_num, val)
    clipboards = [val, get_clipboard(app_num)]

    while not interrupt:
        try:
            for num in ('0', app_num):
                val = get_clipboard(num)
                i = 0 if num == '0' else 1
                if val != clipboards[i]:
                    clipboards[0] = val
                    clipboards[1] = val
                    set_clipboard(app_num if num == '0' else '0', val)
            time.sleep(0.2)
        except:
            if not interrupt:
                traceback.print_exc()
                time.sleep(3)

command = sys.argv[1]
if sys.argv[-1] == '--debug':
    DEBUG = True

if command == 'build':
    app = sys.argv[2]

    exists_pull()
    build(app)

elif command == 'build-all':
    build_all()

elif command == 'update':
    pull()
    build_all()

elif command == 'add':
    app = sys.argv[2]

    exists_pull()
    exists_build(app)
    add(app)

elif command == 'remove':
    app = sys.argv[2]

    remove(app)

elif command == 'run':
    app = sys.argv[2]
    run(app)

else:
    app = sys.argv[1]
    run(app)
<commit_msg>Pull and build on run<commit_after>import os
import shutil
import subprocess
import sys
import random
import threading
import json
import time
import traceback
import signal

USER_HOME_DIR = '~'
HOME_DIR = '~/Docker'
APP_DIR = '~/.local/share/applications'
PULSE_COOKIE_PATH = '~/.config/pulse/cookie'
TMP_DIR = '/tmp'
ROOT_DIR = os.path.dirname(os.path.realpath(__file__))
SCRIPT_PATH = os.path.join(ROOT_DIR, __file__)
PULSE_SERVER = 'unix:/var/run/pulse/native'

# Init
DESKTOP_ENTRY = '''
[Desktop Entry]
Version=1.0
Type=Application
Terminal=false
Name=Docker - %s
Comment=Docker - %s
Exec=%s
Icon=%s
Categories=Other;
'''
USER_HOME_DIR = os.path.expanduser(USER_HOME_DIR)
HOME_DIR = os.path.expanduser(HOME_DIR)
DESKTOP_DIR = os.path.expanduser(APP_DIR)
PULSE_COOKIE_PATH = os.path.expanduser(PULSE_COOKIE_PATH)
TMP_DIR = os.path.expanduser(TMP_DIR)
CONF_DIR = os.path.join(HOME_DIR, '.config')
BASE_CONF_PATH = os.path.join(CONF_DIR, 'base.json')
interrupt = False

def mkdirs(path):
    if not os.path.exists(path):
        os.makedirs(path)

mkdirs(CONF_DIR)

if not os.path.exists(BASE_CONF_PATH):
    shutil.copyfile(
        os.path.join(ROOT_DIR, 'apps', 'base', 'settings.json'),
        BASE_CONF_PATH,
    )

with open(BASE_CONF_PATH, 'r') as conf_file:
    conf_data = json.loads(conf_file.read())

SUDO_DOCKER = conf_data.get('sudo_docker', True)
INCREASE_SHM = conf_data.get('increase_shm', True)
SHARE_CLIPBOARD = conf_data.get('share_clipboard', True)
SHARE_FONTS = conf_data.get('share_fonts', True)
SHARE_THEMES = conf_data.get('share_themes', False)
SHARE_ICONS = conf_data.get('share_icons', False)
SHARE_USER_FONTS = conf_data.get('share_user_fonfs', True)
SHARE_USER_THEMES = conf_data.get('share_user_themes', True)
DPI = conf_data.get('dpi')
DEBUG = False

GPU = conf_data.get('gpu', 'auto')
if GPU == 'auto':
    try:
        subprocess.check_output(['which', 'nvidia-settings'],
            stderr=subprocess.PIPE)
        GPU = 'nvidia'
    except:
        GPU = 'intel'

DEFAULT_WIN_SIZE = conf_data.get('default_win_size', '1024x768')
DEFAULT_VOLUMES = conf_data.get('default_volumes', [])

def kill_process(process):
    terminated = False

    for _ in xrange(200):
        try:
            process.send_signal(signal.SIGINT)
        except OSError as error:
            if error.errno != 3:
                raise
        for _ in xrange(4):
            if process.poll() is not None:
                terminated = True
                break
            time.sleep(0.0025)
        if terminated:
            break

    if not terminated:
        for _ in xrange(10):
            if process.poll() is not None:
                break
            try:
                process.send_signal(signal.SIGKILL)
            except OSError as error:
                if error.errno != 3:
                    raise
            time.sleep(0.01)

def image_exists(image):
    image_id = subprocess.check_output([
        'docker',
        'images',
        '-q',
        image,
    ]).strip()

    return bool(image_id)

def pull():
    subprocess.check_call((['sudo'] if SUDO_DOCKER else []) + [
        'docker',
        'pull',
        'pritunl/archlinux',
    ])

def exists_pull():
    if not image_exists('pritunl/archlinux'):
        pull()

def build(app):
    app = app.split('#')[0]

    if app == 'base-intel' or app == 'base-nvidia':
        image_name = 'base-xorg'
    else:
        image_name = app

    app_dir = os.path.join(ROOT_DIR, 'apps', app)

    subprocess.check_call((['sudo'] if SUDO_DOCKER else []) + [
        'docker',
        'build',
        '--rm',
        '-t', 'dockercask/' + image_name,
        '.',
    ], cwd=app_dir)

    if app == 'base':
        build('base-' + GPU)

def exists_build(app):
    app = app.split('#')[0]

    if not image_exists('dockercask/base'):
        build('base')

    if not image_exists('dockercask/' + app):
        build(app)

def build_all():
    build('base')
    for app in os.listdir(HOME_DIR):
        if not os.path.isdir(os.path.join(HOME_DIR, app)) or \
                app.startswith('.'):
            continue
        build(app)

def add(app):
    app_dir = os.path.join(HOME_DIR, app)
    icon_path = os.path.join(ROOT_DIR, 'apps', app.split('#')[0], 'icon.png')
    app_conf_path = os.path.join(CONF_DIR, app + '.json')
    desktop_entry_path = os.path.join(DESKTOP_DIR,
        'docker-%s.desktop' % app.replace('#', '-'))

    mkdirs(app_dir)

    if not os.path.exists(app_conf_path):
        shutil.copyfile(
            os.path.join(ROOT_DIR, 'apps', app.split('#')[0], 'settings.json'),
            app_conf_path,
        )

    if DEBUG:
        cmd = 'xfce4-terminal --command="python2 %s %s --debug"' % (
            SCRIPT_PATH, app)
    else:
        cmd = 'python2 %s %s' % (SCRIPT_PATH, app)

    formated_app_name = app.replace('#', ' ').replace('-', ' ').split()
    formated_app_name = ' '.join([x.capitalize() for x in formated_app_name])

    with open(desktop_entry_path, 'w') as desktop_file:
        desktop_file.write(DESKTOP_ENTRY % (
            formated_app_name,
            formated_app_name,
            cmd,
            icon_path,
        ))

def remove(app):
    app_dir = os.path.join(HOME_DIR, app)
    app_conf_path = os.path.join(CONF_DIR, app + '.json')
    desktop_entry_path = os.path.join(DESKTOP_DIR,
        'docker-%s.desktop' % app.replace('#', '-'))

    for path in (app_dir, app_conf_path, desktop_entry_path):
        subprocess.check_call([
            'rm',
            '-rf',
            path,
        ])

def run(app):
    try:
        subprocess.check_call([
            'wmctrl',
            '-F',
            '-R',
            app,
        ])
        return
    except subprocess.CalledProcessError:
        pass

    app_dir = os.path.join(HOME_DIR, app)
    app_conf_path = os.path.join(CONF_DIR, app + '.json')
    fonts_dir = os.path.join(USER_HOME_DIR, '.fonts')
    themes_dir = os.path.join(USER_HOME_DIR, '.themes')
    cmd = []
    docker_args = []
    volume_args = []

    if not os.path.exists(app_dir):
        print 'App must be added before running'
        exit(1)

    with open(app_conf_path, 'r') as app_conf_file:
        app_conf_data = json.loads(app_conf_file.read())

    host_x11 = app_conf_data.get('host_x11')
    privileged = app_conf_data.get('privileged')
    increase_shm = app_conf_data.get('increase_shm', INCREASE_SHM)
    dpi = app_conf_data.get('dpi', DPI)

    if DEBUG:
        docker_args.append('-it')
        cmd.append('/bin/bash')

    if privileged:
        docker_args += ['--privileged']

    if increase_shm:
        docker_args += ['--shm-size', '1g']

    for src, dest in DEFAULT_VOLUMES:
        volume_args += [
            '-v', '%s:%s' % (os.path.expanduser(src), dest),
        ]

    if SHARE_FONTS:
        volume_args += [
            '-v', '%s:%s' % ('/usr/share/fonts', '/usr/share/fonts:ro'),
        ]

    if SHARE_ICONS:
        volume_args += [
            '-v', '%s:%s' % ('/usr/share/icons', '/usr/share/icons:ro'),
        ]

    if SHARE_THEMES:
        volume_args += [
            '-v', '%s:%s' % ('/usr/share/themes', '/usr/share/themes:ro'),
        ]

    if SHARE_USER_FONTS:
        volume_args += [
            '-v', '%s:%s' % (fonts_dir, '/home/docker/.fonts:ro'),
        ]

    if SHARE_USER_THEMES:
        volume_args += [
            '-v', '%s:%s' % (themes_dir, '/home/docker/.themes:ro'),
        ]

    mkdirs(fonts_dir)
    mkdirs(themes_dir)

    if host_x11:
        x_cookie = subprocess.check_output(['xauth', 'list', ':0']).split()[-1]
        x_num = '0'
    else:
        x_cookie = subprocess.check_output(['mcookie'])
        x_num = str(random.randint(1000, 32000))
        x_auth_path = os.path.join(TMP_DIR, '.X11-docker-' + x_num)

        with open(x_auth_path, 'w') as _:
            pass

        subprocess.check_call([
            'xauth',
            '-f', x_auth_path,
            'add',
            ':' + x_num,
            'MIT-MAGIC-COOKIE-1',
            x_cookie,
        ])

        subprocess.check_call([
            'xauth',
            'add',
            ':' + x_num,
            'MIT-MAGIC-COOKIE-1',
            x_cookie,
        ])

    x_screen_path = os.path.join(TMP_DIR, '.X11-unix', 'X' + x_num)

    x_proc = None
    docker_id = None

    def clean_up():
        global interrupt
        interrupt = True

        if docker_id:
            try:
                subprocess.check_output([
                    'docker',
                    'rm',
                    '-f',
                    docker_id,
                ], stderr=subprocess.PIPE)
            except:
                pass
        if x_proc:
            kill_process(x_proc)
        if not host_x11:
            try:
                subprocess.check_output([
                    'xauth',
                    'remove',
                    ':' + x_num,
                ], stderr=subprocess.PIPE)
            except:
                pass
            try:
                os.remove(x_auth_path)
            except:
                pass
            try:
                os.remove(x_screen_path)
            except:
                pass

    if not host_x11:
        args = [
            'Xephyr',
            '-auth', x_auth_path,
            '-screen', DEFAULT_WIN_SIZE,
            '-title', app,
            '-br',
            '-resizeable',
            '-no-host-grab',
            '-nolisten', 'tcp',
        ]

        if dpi:
            args += ['-dpi', dpi]

        x_proc = subprocess.Popen(args + [':' + x_num])

        def thread_func():
            try:
                x_proc.wait()
            finally:
                clean_up()

        thread = threading.Thread(target=thread_func)
        thread.start()

    args = (['sudo'] if SUDO_DOCKER else []) + [
        'docker',
        'run',
        '--rm' if DEBUG else '--detach',
    ] + docker_args + [
        '-v', '%s:%s' % (x_screen_path, x_screen_path),
        '-v', '%s:%s' % (PULSE_COOKIE_PATH, '/tmp/.pulse-cookie'),
        '-v', '%s:%s' % (BASE_CONF_PATH, '/base.json:ro'),
        '-v', '%s:%s' % (app_conf_path, '/app.json:ro'),
        '-v', '%s:%s' % (app_dir, '/home/docker/Docker'),
        '-v', '/var/run/user/%s/pulse/native:/var/run/pulse/native' % (
            os.getuid()),
    ] + volume_args + [
        '-u', 'docker',
        '-e', 'HOME=/home/docker',
        '-e', 'DISPLAY=:' + x_num,
        '-e', 'XAUTHORITY=/tmp/.Xauth',
        '-e', 'XCOOKIE=' + x_cookie,
        '-e', 'PULSE_SERVER=' + PULSE_SERVER,
        'dockercask/' + app.split('#')[0],
    ] + cmd

    print ' '.join(args)

    if not host_x11 and SHARE_CLIPBOARD:
        thread = threading.Thread(target=share_clipboard, args=(x_num,))
        thread.daemon = True
        thread.start()

    if DEBUG:
        try:
            subprocess.check_call(args)
        finally:
            clean_up()
    else:
        docker_id = subprocess.check_output(args).strip()

        try:
            subprocess.check_call(['docker', 'wait', docker_id])
        finally:
            clean_up()

def set_clipboard(num, val):
    proc = subprocess.Popen(
        ['xsel', '--display', ':' + num, '-b', '-i', '-t', '500'],
        stdin=subprocess.PIPE,
    )
    proc.communicate(input=val)
    proc.wait()

def get_clipboard(num):
    return subprocess.check_output(
        ['xsel', '--display', ':' + num, '-b', '-o', '-t', '500'],
    )

def share_clipboard(app_num):
    time.sleep(1)

    val = get_clipboard('0')
    set_clipboard(app_num, val)
    clipboards = [val, get_clipboard(app_num)]

    while not interrupt:
        try:
            for num in ('0', app_num):
                val = get_clipboard(num)
                i = 0 if num == '0' else 1
                if val != clipboards[i]:
                    clipboards[0] = val
                    clipboards[1] = val
                    set_clipboard(app_num if num == '0' else '0', val)
            time.sleep(0.2)
        except:
            if not interrupt:
                traceback.print_exc()
                time.sleep(3)

command = sys.argv[1]
if sys.argv[-1] == '--debug':
    DEBUG = True

if command == 'build':
    app = sys.argv[2]

    exists_pull()
    build(app)

elif command == 'build-all':
    build_all()

elif command == 'update':
    pull()
    build_all()

elif command == 'add':
    app = sys.argv[2]

    exists_pull()
    exists_build(app)
    add(app)

elif command == 'remove':
    app = sys.argv[2]

    remove(app)

elif command == 'run':
    app = sys.argv[2]

    exists_pull()
    exists_build(app)
    run(app)

else:
    app = sys.argv[1]

    exists_pull()
    exists_build(app)
    run(app)
<|endoftext|>


<commit_before>import concurrent.futures
import http
import http.client
import socket
import sys
import wsgiref.simple_server

class ThreadPoolServer(wsgiref.simple_server.WSGIServer):
    def __pre_init(self, max_threads):
        '''
        This must be called, by a deriving class, before __init__ is called.

        This is because the thread pool has to be set up before requests are
        processed, and because overriding __init__ is problematic while also
        changing its signature, *and* coöperating with IPv64Server.
        '''
        if sys.version_info.major <= 3 and sys.version_info.minor < 5:
            if max_threads is None:
                max_threads = 4
        self.__ex = concurrent.futures.ThreadPoolExecutor(max_threads)

    def process_request(self, request, client_address):
        self.__ex.submit(self.__process_request_thread, request, client_address)

    def __process_request_thread(self, request, client_address):
        '''
        Taken from socketserver.ThreadingMixIn
        '''
        try:
            self.finish_request(request, client_address)
            self.shutdown_request(request)
        except Exception:
            self.handle_error(request, client_address)
            self.shutdown_request(request)

    def server_close(self):
        super().server_close()
        self.__ex.shutdown()

class InstantShutdownServer(wsgiref.simple_server.WSGIServer):
    '''
    Connecting to the underlying SocketServer's listening socket will wake it
    up. It will then immediately check the shutdown flag, rather than waiting
    for the poll_interval.
    '''
    def shutdown(self):
        with socket.socket(self.socket.family) as s:
            s.setblocking(0)
            try:
                s.connect(self.socket.getsockname())
            except BlockingIOError:
                pass
            # Now the thread running serve_forever is waiting for the client
            # to send a request...
            super().shutdown()
        # Now our socket is closed, the thread running serve_forever will check
        # __shutdown_request and return.

class IPv64Server(wsgiref.simple_server.WSGIServer):
    def __pre_init(self, server_address, bind_v6only):
        '''
        This must be called, by a deriving class, before __init__ is called.

        This is because TCPServer.__init__ uses self.address_family to create
        the socket.
        '''
        self.address_family = socket.AF_INET6 if server_address.version == 6 else socket.AF_INET
        self.__bind_v6only = bind_v6only

    def server_bind(self):
        self.socket.setsockopt(socket.IPPROTO_IP, 15, 1) # IP_FREEBIND
        if self.__bind_v6only is not None and self.address_family == socket.AF_INET6:
            self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, self.__bind_v6only)
        super().server_bind()

class HealthCheckServer(wsgiref.simple_server.WSGIServer):
    def healthy(self):
        c = http.client.HTTPConnection(self.server_address[0], self.server_address[1], timeout=5)
        try:
            c.request('GET', '/')
        except http.client.HTTPException:
            return False
        with c.getresponse() as r:
            if r.status != 200:
                return False
        return True

class SilentRequestHandler(wsgiref.simple_server.WSGIRequestHandler):
    def log_request(self, code='-', message='-'):
        if isinstance(code, http.HTTPStatus) and code.value < 400:
            return
        elif code[0] < '4':
            return
        super().log_request(code, message)

class Server(HealthCheckServer, IPv64Server, InstantShutdownServer, ThreadPoolServer):
    '''
    A WSGIServer that works with IPv6, and processes requests concurrently.

    server_address[0] must be an ipaddress.ip_address, as opposed to the normal string.
    '''
    def __init__(self, server_address, RequestHandlerClass, max_threads, bind_v6only, bind_and_activate=True):
        self._IPv64Server__pre_init(server_address[0], bind_v6only)
        self._ThreadPoolServer__pre_init(max_threads)
        super().__init__((str(server_address[0]), server_address[1]), RequestHandlerClass, bind_and_activate)
<commit_msg>Use HTTPStatus instead of magic number<commit_after>import concurrent.futures
import http
import http.client
import socket
import sys
import wsgiref.simple_server

class ThreadPoolServer(wsgiref.simple_server.WSGIServer):
    def __pre_init(self, max_threads):
        '''
        This must be called, by a deriving class, before __init__ is called.

        This is because the thread pool has to be set up before requests are
        processed, and because overriding __init__ is problematic while also
        changing its signature, *and* coöperating with IPv64Server.
        '''
        if sys.version_info.major <= 3 and sys.version_info.minor < 5:
            if max_threads is None:
                max_threads = 4
        self.__ex = concurrent.futures.ThreadPoolExecutor(max_threads)

    def process_request(self, request, client_address):
        self.__ex.submit(self.__process_request_thread, request, client_address)

    def __process_request_thread(self, request, client_address):
        '''
        Taken from socketserver.ThreadingMixIn
        '''
        try:
            self.finish_request(request, client_address)
            self.shutdown_request(request)
        except Exception:
            self.handle_error(request, client_address)
            self.shutdown_request(request)

    def server_close(self):
        super().server_close()
        self.__ex.shutdown()

class InstantShutdownServer(wsgiref.simple_server.WSGIServer):
    '''
    Connecting to the underlying SocketServer's listening socket will wake it
    up. It will then immediately check the shutdown flag, rather than waiting
    for the poll_interval.
    '''
    def shutdown(self):
        with socket.socket(self.socket.family) as s:
            s.setblocking(0)
            try:
                s.connect(self.socket.getsockname())
            except BlockingIOError:
                pass
            # Now the thread running serve_forever is waiting for the client
            # to send a request...
            super().shutdown()
        # Now our socket is closed, the thread running serve_forever will check
        # __shutdown_request and return.

class IPv64Server(wsgiref.simple_server.WSGIServer):
    def __pre_init(self, server_address, bind_v6only):
        '''
        This must be called, by a deriving class, before __init__ is called.

        This is because TCPServer.__init__ uses self.address_family to create
        the socket.
        '''
        self.address_family = socket.AF_INET6 if server_address.version == 6 else socket.AF_INET
        self.__bind_v6only = bind_v6only

    def server_bind(self):
        self.socket.setsockopt(socket.IPPROTO_IP, 15, 1) # IP_FREEBIND
        if self.__bind_v6only is not None and self.address_family == socket.AF_INET6:
            self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, self.__bind_v6only)
        super().server_bind()

class HealthCheckServer(wsgiref.simple_server.WSGIServer):
    def healthy(self):
        c = http.client.HTTPConnection(self.server_address[0], self.server_address[1], timeout=5)
        try:
            c.request('GET', '/')
        except http.client.HTTPException:
            return False
        with c.getresponse() as r:
            if r.status != http.HTTPStatus.OK:
                return False
        return True

class SilentRequestHandler(wsgiref.simple_server.WSGIRequestHandler):
    def log_request(self, code='-', message='-'):
        if isinstance(code, http.HTTPStatus) and code.value < 400:
            return
        elif code[0] < '4':
            return
        super().log_request(code, message)

class Server(HealthCheckServer, IPv64Server, InstantShutdownServer, ThreadPoolServer):
    '''
    A WSGIServer that works with IPv6, and processes requests concurrently.

    server_address[0] must be an ipaddress.ip_address, as opposed to the normal string.
    '''
    def __init__(self, server_address, RequestHandlerClass, max_threads, bind_v6only, bind_and_activate=True):
        self._IPv64Server__pre_init(server_address[0], bind_v6only)
        self._ThreadPoolServer__pre_init(max_threads)
        super().__init__((str(server_address[0]), server_address[1]), RequestHandlerClass, bind_and_activate)
<|endoftext|>


<commit_before>import concurrent.futures
import http
import http.client
import ipaddress
import socket
import socketserver
import sys
import wsgiref.simple_server

class ThreadPoolServer(socketserver.TCPServer):
    def __pre_init(self, max_threads):
        '''
        This must be called, by a deriving class, before __init__ is called.

        This is because the thread pool has to be set up before requests are
        processed, and because overriding __init__ is problematic while also
        changing its signature, *and* coöperating with IPv64Server.
        '''
        if sys.version_info.major <= 3 and sys.version_info.minor < 5:
            if max_threads is None:
                max_threads = 4
        self.__ex = concurrent.futures.ThreadPoolExecutor(max_threads)

    def process_request(self, request, client_address):
        self.__ex.submit(self.__process_request_thread, request, client_address)

    def __process_request_thread(self, request, client_address):
        '''
        Taken from socketserver.ThreadingMixIn
        '''
        try:
            self.finish_request(request, client_address)
            self.shutdown_request(request)
        except Exception:
            self.handle_error(request, client_address)
            self.shutdown_request(request)

    def server_close(self):
        super().server_close()
        self.__ex.shutdown()

class InstantShutdownServer(socketserver.TCPServer):
    '''
    Connecting to the underlying SocketServer's listening socket will wake it
    up. It will then immediately check the shutdown flag, rather than waiting
    for the poll_interval.
    '''
    def send_stop(self):
        # Unfortunately there is no public API for setting this flag, and
        # calling super().shutdown() here would block us
        self._BaseServer__shutdown_request = True
        # Now connect to the server to cause it to wake up
        with socket.socket(self.socket.family) as s:
            s.setblocking(0)
            try:
                s.connect(self.socket.getsockname())
            except BlockingIOError:
                pass

class IPv64Server(socketserver.TCPServer):
    def __pre_init(self, server_address, bind_v6only):
        '''
        This must be called, by a deriving class, before __init__ is called.

        This is because TCPServer.__init__ uses self.address_family to create
        the socket.
        '''
        ip = ipaddress.ip_address(server_address[0])
        self.address_family = socket.AF_INET6 if ip.version == 6 else socket.AF_INET
        self.__bind_v6only = bind_v6only

    def server_bind(self):
        self.socket.setsockopt(socket.IPPROTO_IP, 15, 1) # IP_FREEBIND
        if self.__bind_v6only is not None and self.address_family == socket.AF_INET6:
            self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, self.__bind_v6only)
        super().server_bind()

class HealthCheckServer(wsgiref.simple_server.WSGIServer):
    def healthy(self):
        c = http.client.HTTPConnection(self.server_address[0], self.server_address[1], timeout=5)
        try:
            c.request('GET', '/')
        except http.client.HTTPException:
            return False
        with c.getresponse() as r:
            if r.status != http.HTTPStatus.OK:
                return False
        return True

class SilentRequestHandler(wsgiref.simple_server.WSGIRequestHandler):
    def log_request(self, code='-', message='-'):
        if isinstance(code, http.HTTPStatus) and code.value < 400:
            return
        elif code[0] < '4':
            return
        super().log_request(code, message)

class Server(HealthCheckServer, IPv64Server, InstantShutdownServer, ThreadPoolServer, wsgiref.simple_server.WSGIServer):
    def __init__(self, server_address, *, max_threads, bind_v6only, bind_and_activate=True):
        self._IPv64Server__pre_init(server_address, bind_v6only)
        self._ThreadPoolServer__pre_init(max_threads)
        super().__init__(server_address, SilentRequestHandler, bind_and_activate)
<commit_msg>Move common case of log_request filter test to be first, and explain when it is called with a str vs an HTTPStatus<commit_after>import concurrent.futures
import http
import http.client
import ipaddress
import socket
import socketserver
import sys
import wsgiref.simple_server

class ThreadPoolServer(socketserver.TCPServer):
    def __pre_init(self, max_threads):
        '''
        This must be called, by a deriving class, before __init__ is called.

        This is because the thread pool has to be set up before requests are
        processed, and because overriding __init__ is problematic while also
        changing its signature, *and* coöperating with IPv64Server.
        '''
        if sys.version_info.major <= 3 and sys.version_info.minor < 5:
            if max_threads is None:
                max_threads = 4
        self.__ex = concurrent.futures.ThreadPoolExecutor(max_threads)

    def process_request(self, request, client_address):
        self.__ex.submit(self.__process_request_thread, request, client_address)

    def __process_request_thread(self, request, client_address):
        '''
        Taken from socketserver.ThreadingMixIn
        '''
        try:
            self.finish_request(request, client_address)
            self.shutdown_request(request)
        except Exception:
            self.handle_error(request, client_address)
            self.shutdown_request(request)

    def server_close(self):
        super().server_close()
        self.__ex.shutdown()

class InstantShutdownServer(socketserver.TCPServer):
    '''
    Connecting to the underlying SocketServer's listening socket will wake it
    up. It will then immediately check the shutdown flag, rather than waiting
    for the poll_interval.
    '''
    def send_stop(self):
        # Unfortunately there is no public API for setting this flag, and
        # calling super().shutdown() here would block us
        self._BaseServer__shutdown_request = True
        # Now connect to the server to cause it to wake up
        with socket.socket(self.socket.family) as s:
            s.setblocking(0)
            try:
                s.connect(self.socket.getsockname())
            except BlockingIOError:
                pass

class IPv64Server(socketserver.TCPServer):
    def __pre_init(self, server_address, bind_v6only):
        '''
        This must be called, by a deriving class, before __init__ is called.

        This is because TCPServer.__init__ uses self.address_family to create
        the socket.
        '''
        ip = ipaddress.ip_address(server_address[0])
        self.address_family = socket.AF_INET6 if ip.version == 6 else socket.AF_INET
        self.__bind_v6only = bind_v6only

    def server_bind(self):
        self.socket.setsockopt(socket.IPPROTO_IP, 15, 1) # IP_FREEBIND
        if self.__bind_v6only is not None and self.address_family == socket.AF_INET6:
            self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, self.__bind_v6only)
        super().server_bind()

class HealthCheckServer(wsgiref.simple_server.WSGIServer):
    def healthy(self):
        c = http.client.HTTPConnection(self.server_address[0], self.server_address[1], timeout=5)
        try:
            c.request('GET', '/')
        except http.client.HTTPException:
            return False
        with c.getresponse() as r:
            if r.status != http.HTTPStatus.OK:
                return False
        return True

class SilentRequestHandler(wsgiref.simple_server.WSGIRequestHandler):
    def log_request(self, code='-', message='-'):
        if isinstance(code, str) and code[0] < '4':
            # WSGIRequestHandler always calls this with a string.
            return
        elif isinstance(code, http.HTTPStatus) and code.value < 400:
            # But a bad request is handled by BaseHTTPRequestHandler, which
            # uses an HTTPStatus
            return
        super().log_request(code, message)

class Server(HealthCheckServer, IPv64Server, InstantShutdownServer, ThreadPoolServer, wsgiref.simple_server.WSGIServer):
    def __init__(self, server_address, *, max_threads, bind_v6only, bind_and_activate=True):
        self._IPv64Server__pre_init(server_address, bind_v6only)
        self._ThreadPoolServer__pre_init(max_threads)
        super().__init__(server_address, SilentRequestHandler, bind_and_activate)
<|endoftext|>


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

version = '0.1.9'
<commit_msg>0.1.10<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

version = '0.1.10'
<|endoftext|>


<commit_before><commit_msg>mods for interleaved idle identity in RB<commit_after><|endoftext|>