<commit_before>import json
import logging
from collections import OrderedDict
from io import BytesIO

from reportlab.lib.enums import TA_CENTER
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import ListStyle, ParagraphStyle, getSampleStyleSheet
from reportlab.pdfgen import canvas
from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer

from django.conf import settings
from django.utils import timezone
from django.utils.html import conditional_escape

from callisto_core.utils import api

logger = logging.getLogger(__name__)

def report_as_pdf(report, data, recipient):
    return PDFFullReport(
        report=report,
        report_data=data,
    ).generate_pdf_report(
        recipient=recipient,
        report_id=report.id,
    )

class MatchReportContent(object):
    '''
        Class to structure contact information collected
        from match submission form for report
    '''

    # This constructor is called with keyword arguments populated by
    # encrypted data. Existing arguments should not be removed or renamed,
    # and new arguments must have default values.
    def __init__(
        self, identifier, perp_name, email, phone,
        contact_name=None, voicemail=None, notes=None,
    ):
        self.identifier = conditional_escape(identifier)
        self.perp_name = conditional_escape(perp_name)
        self.contact_name = conditional_escape(contact_name)
        self.email = conditional_escape(email)
        self.phone = conditional_escape(phone)
        self.voicemail = conditional_escape(voicemail)
        self.notes = conditional_escape(notes)

    @classmethod
    def from_form(cls, form):
        return cls(
            identifier=form.cleaned_data.get('identifier'),
            perp_name=form.cleaned_data.get('perp_name'),
            contact_name=form.instance.report.contact_name,
            email=form.instance.report.contact_email,
            phone=form.instance.report.contact_phone,
            voicemail=form.instance.report.contact_voicemail,
            notes=form.instance.report.contact_notes,
        )

class NumberedCanvas(canvas.Canvas):

    def __init__(self, *args, **kwargs):
        canvas.Canvas.__init__(self, *args, **kwargs)
        self._saved_page_states = []

    def showPage(self):
        self._saved_page_states.append(dict(self.__dict__))
        self._startPage()

    def save(self):
        '''add page info to each page (page x of y)'''
        num_pages = len(self._saved_page_states)
        for state in self._saved_page_states:
            self.__dict__.update(state)
            self.draw_page_number(num_pages)
            canvas.Canvas.showPage(self)
        canvas.Canvas.save(self)

    def draw_page_number(self, page_count):
        width, height = letter
        margin = 0.66 * 72
        self.setFillColor('gray')
        self.drawRightString(
            width - margin, margin, "Page %d of %d" %
            (self._pageNumber, page_count))

class PDFReport(object):

    unselected = u'\u2610'
    selected = u'\u2717'
    free_text = u'\u2756'
    no_bullet = ' '

    report_title = "Report"

    def __init__(self):
        self.styles = self.set_up_styles()
        self.notes_style = self.styles['Notes']
        self.answers_style = self.styles['Answers']
        self.answers_list_style = self.styles['AnswersList']
        self.body_style = self.styles['Normal']
        self.section_title_style = self.styles['Heading4']
        self.report_title_style = self.styles['Heading3']
        self.pdf_elements = []

    def set_up_styles(self, *args, **kwargs):
        '''
        Helper formatting methods assume the following styles have been defined:
            * Normal: format_question, get_metadata_page
            * Notes: get_metadata_page
            * Answers: format_answer, get_metadata_page
            * AnswersList: format_answer_list, get_metadata_page
            * Heading4: get_metadata_page, PDFFullReport.generate_pdf_report, PDFMatchReport.generate_match_report
            * Heading3: get_metadata_page, PDFMatchReport.generate_match_report
        If you override this method but don't define these styles, you will need to also override init and potentially
        the corresponding formatting methods.
        :return: a ReportLab stylesheet
        '''
        def get_font_location(filename):
            return settings.APPS_DIR.path('fonts')(filename)

        styles = getSampleStyleSheet()
        headline_style = styles["Heading1"]
        headline_style.alignment = TA_CENTER
        headline_style.fontSize = 48

        subtitle_style = styles["Heading2"]
        subtitle_style.fontSize = 24
        subtitle_style.leading = 26
        subtitle_style.alignment = TA_CENTER

        report_title_style = styles["Heading3"]
        report_title_style.fontSize = 28
        report_title_style.leading = 36
        report_title_style.spaceBefore = 0
        report_title_style.alignment = TA_CENTER

        section_title_style = styles["Heading4"]
        section_title_style.fontSize = 18
        section_title_style.leading = 22
        section_title_style.spaceBefore = 20

        body_style = styles["Normal"]
        body_style.fontSize = 14
        body_style.leading = 17
        body_style.leftIndent = 25

        styles.add(ParagraphStyle(
            name='Notes',
            parent=styles['Normal'],
            fontSize=10,
            leading=14,
            leftIndent=45,
            alias='notes',
        ))
        styles.add(ListStyle(
            name='AnswersList',
        ))
        styles.add(ParagraphStyle(
            name='Answers',
            parent=styles['Normal'],
            leftIndent=0,
            alias='answers',
        ))

        return styles

    # FORMATTING HELPERS

    def add_question(self, question):
        self.pdf_elements.append(
            Paragraph(question, self.body_style),
        )
        self.pdf_elements.append(Spacer(1, 4))

    def add_answer_list(self, answers):
        for answer in answers:
            self.pdf_elements.append(
                Paragraph(answer, self.notes_style),
            )
            self.pdf_elements.append(Spacer(1, 1))

    def render_question(self, question, answers):
        self.add_question(question)
        self.add_answer_list(answers)

    def get_header_footer(self, recipient):
        def func(canvas, doc):
            width, height = letter
            margin = 0.66 * 72
            canvas.saveState()
            canvas.setFillColor('gray')
            canvas.drawString(margin, height - margin, "CONFIDENTIAL")
            canvas.drawRightString(
                width - margin, height - margin, str(timezone.now()))
            canvas.drawString(
                margin, margin,
                f"Intended for: Title IX Coordinator {recipient}",
            )
            canvas.restoreState()
        return func

    @staticmethod
    def get_user_identifier(user):
        try:
            return user.email or user.username
        except AttributeError:
            return 'Anonymous User'

class PDFFullReport(PDFReport):

    def __init__(self, report, report_data):
        super().__init__()
        self.report = report
        self.report_data = report_data

    def get_metadata_page(self, recipient):
        MetadataPage = []
        MetadataPage.append(Paragraph(
            api.NotificationApi.report_title,
            self.report_title_style,
        ))

        MetadataPage.append(Paragraph("Overview", self.section_title_style))

        overview_body = "Reported by: {0}
".format(
            self.get_user_identifier(self.report.owner),
        )
        overview_body = overview_body + '''
            Submitted on: {0}

        '''.format(
            self.report.submitted_to_school,
        )
        overview_body = overview_body + '''
            Record Created: {0}

        '''.format(self.report.added.strftime("%Y-%m-%d %H:%M"))

        MetadataPage.append(Paragraph(overview_body, self.body_style))

        MetadataPage.append(Paragraph(
            "Contact Preferences",
            self.section_title_style,
        ))

        contact_body = '''
            Name: {0}

            Phone: {1}

            Voicemail preferences: {2}

            Email: {3}

            Notes on preferred contact time of day, gender of admin, etc.:
        '''.format(
            self.report.contact_name or "None provided",
            self.report.contact_phone,
            self.report.contact_voicemail or "None provided",
            self.report.contact_email,
        )

        MetadataPage.append(Paragraph(
            contact_body,
            self.body_style,
        ))
        MetadataPage.append(Paragraph(
            self.report.contact_notes or "None provided",
            self.notes_style,
        ))

        return MetadataPage

    def generate_pdf_report(self, report_id, recipient):
        # PREPARE PDF
        report_buffer = BytesIO()
        doc = SimpleDocTemplate(
            report_buffer,
            pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )

        # COVER PAGE
        self.pdf_elements.extend(
            api.NotificationApi.get_cover_page(
                report_id=report_id,
                recipient=recipient,
            ),
        )

        # METADATA PAGE
        self.pdf_elements.extend(self.get_metadata_page(recipient))

        # REPORT
        self.pdf_elements.extend(
            [Paragraph("Report Questions", self.section_title_style)],
        )
        for item in self.report_data:
            question, answers = item.popitem()
            self.render_question(question, answers)

        doc.build(
            self.pdf_elements,
            onFirstPage=self.get_header_footer(recipient),
            onLaterPages=self.get_header_footer(recipient),
            canvasmaker=NumberedCanvas,
        )
        result = report_buffer.getvalue()
        report_buffer.close()
        return result

class PDFMatchReport(PDFReport):

    report_title = "Match Report"

    def __init__(self, matches, identifier):
        super(PDFMatchReport, self).__init__()
        self.matches = sorted(
            matches,
            key=lambda m: m.added.strftime("%Y-%m-%d %H:%M"),
            reverse=True)
        self.identifier = identifier

    def generate_match_report(self, report_id, recipient):
        '''
        Generates PDF report about a discovered match.

        Args:
          report_id (str): id used to uniquely identify this report to receiving authority

        Returns:
          bytes: a PDF with the submitted perp information & contact information of the reporters for this match

        '''
        matches_with_reports = [
            (match,
             MatchReportContent(
                 **json.loads(match.get_match(self.identifier))))
            for match in self.matches]

        buffer = BytesIO()
        doc = SimpleDocTemplate(
            buffer, pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )
        # COVER PAGE
        self.pdf_elements.extend(
            api.NotificationApi.get_cover_page(
                report_id=report_id,
                recipient=recipient,
            ),
        )

        # MATCH REPORTS
        self.pdf_elements.append(
            Paragraph(
                api.NotificationApi.report_title,
                self.report_title_style,
            )
        )

        # perpetrator info
        self.pdf_elements.append(
            Paragraph(
                "Perpetrator",
                self.section_title_style,
            )
        )
        names = ', '.join(OrderedDict.fromkeys([report.perp_name.strip(
        ) for _, report in matches_with_reports if report.perp_name]))
        if len(names) < 1:
            names = 'None provided'
        perp_info = ('Name(s): ' + names) + '
' + \
            "Matching identifier: " + self.identifier
        self.pdf_elements.append(Paragraph(perp_info, self.body_style))

        # reporter info
        for idx, (match_report, match_report_content) in enumerate(
                matches_with_reports):
            self.pdf_elements.append(
                Paragraph(
                    "Report " + str(idx + 1),
                    self.section_title_style,
                )
            )

            report = match_report.report
            user = report.owner

            if report.submitted_to_school:
                sent_report = report.sentfullreport_set.first()
                report_id = sent_report.get_report_id() if sent_report else "Not found"
                is_submitted = '''
                    Yes
                    Submitted to school on: {0}
                    Submitted report ID: {1}
                '''.format(
                    report.submitted_to_school,
                    report_id,
                )
            else:
                is_submitted = "No"

            overview_body = '''
                Perpetrator name given: {0}
                Reported by: {1}
                Submitted to matching on: {2}
                Record created: {3}
                Full record submitted? {4}
            '''.format(
                match_report_content.perp_name or "None provided",
                self.get_user_identifier(user),
                match_report.added.strftime("%Y-%m-%d %H:%M"),
                report.added.strftime("%Y-%m-%d %H:%M"),
                is_submitted,
            ).replace('\n', '
\n')

            self.pdf_elements.append(Paragraph(overview_body, self.body_style))

            contact_body = '''
Name: {0}
                              Phone: {1}
                              Voicemail preferences: {2}
                              Email: {3}
                              Notes on preferred contact time of day, gender of admin, etc.:'''.format(
                match_report_content.contact_name or "None provided",
                match_report_content.phone,
                match_report_content.voicemail or "None provided",
                match_report_content.email).replace('\n', '
\n')

            self.pdf_elements.append(Paragraph(contact_body, self.body_style))
            self.pdf_elements.append(Paragraph(
                match_report_content.notes or "None provided",
                self.notes_style,
            ))

        doc.build(
            self.pdf_elements,
            onFirstPage=self.get_header_footer(recipient),
            onLaterPages=self.get_header_footer(recipient),
            canvasmaker=NumberedCanvas,
        )
        result = buffer.getvalue()
        buffer.close()
        return result
<commit_msg>PDFUserReviewReport<commit_after>import json
import logging
from collections import OrderedDict
from io import BytesIO

from reportlab.lib.enums import TA_CENTER
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import ListStyle, ParagraphStyle, getSampleStyleSheet
from reportlab.pdfgen import canvas
from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer

from django.conf import settings
from django.utils import timezone
from django.utils.html import conditional_escape

from callisto_core.utils import api

logger = logging.getLogger(__name__)

def report_as_pdf(report, data, recipient):
    return PDFFullReport(
        report=report,
        report_data=data,
    ).generate_pdf_report(
        recipient=recipient,
        report_id=report.id,
    )

class MatchReportContent(object):
    '''
        Class to structure contact information collected
        from match submission form for report
    '''

    # This constructor is called with keyword arguments populated by
    # encrypted data. Existing arguments should not be removed or renamed,
    # and new arguments must have default values.
    def __init__(
        self, identifier, perp_name, email, phone,
        contact_name=None, voicemail=None, notes=None,
    ):
        self.identifier = conditional_escape(identifier)
        self.perp_name = conditional_escape(perp_name)
        self.contact_name = conditional_escape(contact_name)
        self.email = conditional_escape(email)
        self.phone = conditional_escape(phone)
        self.voicemail = conditional_escape(voicemail)
        self.notes = conditional_escape(notes)

    @classmethod
    def from_form(cls, form):
        return cls(
            identifier=form.cleaned_data.get('identifier'),
            perp_name=form.cleaned_data.get('perp_name'),
            contact_name=form.instance.report.contact_name,
            email=form.instance.report.contact_email,
            phone=form.instance.report.contact_phone,
            voicemail=form.instance.report.contact_voicemail,
            notes=form.instance.report.contact_notes,
        )

class NumberedCanvas(canvas.Canvas):

    def __init__(self, *args, **kwargs):
        canvas.Canvas.__init__(self, *args, **kwargs)
        self._saved_page_states = []

    def showPage(self):
        self._saved_page_states.append(dict(self.__dict__))
        self._startPage()

    def save(self):
        '''add page info to each page (page x of y)'''
        num_pages = len(self._saved_page_states)
        for state in self._saved_page_states:
            self.__dict__.update(state)
            self.draw_page_number(num_pages)
            canvas.Canvas.showPage(self)
        canvas.Canvas.save(self)

    def draw_page_number(self, page_count):
        width, height = letter
        margin = 0.66 * 72
        self.setFillColor('gray')
        self.drawRightString(
            width - margin, margin, "Page %d of %d" %
            (self._pageNumber, page_count))

class PDFReport(object):

    unselected = u'\u2610'
    selected = u'\u2717'
    free_text = u'\u2756'
    no_bullet = ' '

    report_title = "Report"

    def __init__(self):
        self.styles = self.set_up_styles()
        self.notes_style = self.styles['Notes']
        self.answers_style = self.styles['Answers']
        self.answers_list_style = self.styles['AnswersList']
        self.body_style = self.styles['Normal']
        self.section_title_style = self.styles['Heading4']
        self.report_title_style = self.styles['Heading3']
        self.pdf_elements = []

    def set_up_styles(self, *args, **kwargs):
        '''
        Helper formatting methods assume the following styles have been defined:
            * Normal: format_question, get_metadata_page
            * Notes: get_metadata_page
            * Answers: format_answer, get_metadata_page
            * AnswersList: format_answer_list, get_metadata_page
            * Heading4: get_metadata_page, PDFFullReport.generate_pdf_report, PDFMatchReport.generate_match_report
            * Heading3: get_metadata_page, PDFMatchReport.generate_match_report
        If you override this method but don't define these styles, you will need to also override init and potentially
        the corresponding formatting methods.
        :return: a ReportLab stylesheet
        '''
        def get_font_location(filename):
            return settings.APPS_DIR.path('fonts')(filename)

        styles = getSampleStyleSheet()
        headline_style = styles["Heading1"]
        headline_style.alignment = TA_CENTER
        headline_style.fontSize = 48

        subtitle_style = styles["Heading2"]
        subtitle_style.fontSize = 24
        subtitle_style.leading = 26
        subtitle_style.alignment = TA_CENTER

        report_title_style = styles["Heading3"]
        report_title_style.fontSize = 28
        report_title_style.leading = 36
        report_title_style.spaceBefore = 0
        report_title_style.alignment = TA_CENTER

        section_title_style = styles["Heading4"]
        section_title_style.fontSize = 18
        section_title_style.leading = 22
        section_title_style.spaceBefore = 20

        body_style = styles["Normal"]
        body_style.fontSize = 14
        body_style.leading = 17
        body_style.leftIndent = 25

        styles.add(ParagraphStyle(
            name='Notes',
            parent=styles['Normal'],
            fontSize=10,
            leading=14,
            leftIndent=45,
            alias='notes',
        ))
        styles.add(ListStyle(
            name='AnswersList',
        ))
        styles.add(ParagraphStyle(
            name='Answers',
            parent=styles['Normal'],
            leftIndent=0,
            alias='answers',
        ))

        return styles

    # FORMATTING HELPERS

    def add_question(self, question):
        self.pdf_elements.append(
            Paragraph(question, self.body_style),
        )
        self.pdf_elements.append(Spacer(1, 4))

    def add_answer_list(self, answers):
        for answer in answers:
            self.pdf_elements.append(
                Paragraph(answer, self.notes_style),
            )
            self.pdf_elements.append(Spacer(1, 1))

    def render_question(self, question, answers):
        self.add_question(question)
        self.add_answer_list(answers)

    def get_header_footer(self, recipient):
        def func(canvas, doc):
            width, height = letter
            margin = 0.66 * 72
            canvas.saveState()
            canvas.setFillColor('gray')
            canvas.drawString(margin, height - margin, "CONFIDENTIAL")
            canvas.drawRightString(
                width - margin, height - margin, str(timezone.now()))
            canvas.drawString(
                margin, margin,
                f"Intended for: Title IX Coordinator {recipient}",
            )
            canvas.restoreState()
        return func

    @staticmethod
    def get_user_identifier(user):
        try:
            return user.email or user.username
        except AttributeError:
            return 'Anonymous User'

class PDFUserReviewReport(PDFReport):

    def __init__(self, pdf_input_data):
        super().__init__()
        self.pdf_input_data = pdf_input_data

class PDFFullReport(PDFReport):

    def __init__(self, report, report_data):
        super().__init__()
        self.report = report
        self.report_data = report_data

    def get_metadata_page(self, recipient):
        MetadataPage = []
        MetadataPage.append(Paragraph(
            api.NotificationApi.report_title,
            self.report_title_style,
        ))

        MetadataPage.append(Paragraph("Overview", self.section_title_style))

        overview_body = "Reported by: {0}
".format(
            self.get_user_identifier(self.report.owner),
        )
        overview_body = overview_body + '''
            Submitted on: {0}

        '''.format(
            self.report.submitted_to_school,
        )
        overview_body = overview_body + '''
            Record Created: {0}

        '''.format(self.report.added.strftime("%Y-%m-%d %H:%M"))

        MetadataPage.append(Paragraph(overview_body, self.body_style))

        MetadataPage.append(Paragraph(
            "Contact Preferences",
            self.section_title_style,
        ))

        contact_body = '''
            Name: {0}

            Phone: {1}

            Voicemail preferences: {2}

            Email: {3}

            Notes on preferred contact time of day, gender of admin, etc.:
        '''.format(
            self.report.contact_name or "None provided",
            self.report.contact_phone,
            self.report.contact_voicemail or "None provided",
            self.report.contact_email,
        )

        MetadataPage.append(Paragraph(
            contact_body,
            self.body_style,
        ))
        MetadataPage.append(Paragraph(
            self.report.contact_notes or "None provided",
            self.notes_style,
        ))

        return MetadataPage

    def generate_pdf_report(self, report_id, recipient):
        # PREPARE PDF
        report_buffer = BytesIO()
        doc = SimpleDocTemplate(
            report_buffer,
            pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )

        # COVER PAGE
        self.pdf_elements.extend(
            api.NotificationApi.get_cover_page(
                report_id=report_id,
                recipient=recipient,
            ),
        )

        # METADATA PAGE
        self.pdf_elements.extend(self.get_metadata_page(recipient))

        # REPORT
        self.pdf_elements.extend(
            [Paragraph("Report Questions", self.section_title_style)],
        )
        for item in self.report_data:
            question, answers = item.popitem()
            self.render_question(question, answers)

        doc.build(
            self.pdf_elements,
            onFirstPage=self.get_header_footer(recipient),
            onLaterPages=self.get_header_footer(recipient),
            canvasmaker=NumberedCanvas,
        )
        result = report_buffer.getvalue()
        report_buffer.close()
        return result

class PDFMatchReport(PDFReport):

    report_title = "Match Report"

    def __init__(self, matches, identifier):
        super(PDFMatchReport, self).__init__()
        self.matches = sorted(
            matches,
            key=lambda m: m.added.strftime("%Y-%m-%d %H:%M"),
            reverse=True)
        self.identifier = identifier

    def generate_match_report(self, report_id, recipient):
        '''
        Generates PDF report about a discovered match.

        Args:
          report_id (str): id used to uniquely identify this report to receiving authority

        Returns:
          bytes: a PDF with the submitted perp information & contact information of the reporters for this match

        '''
        matches_with_reports = [
            (match,
             MatchReportContent(
                 **json.loads(match.get_match(self.identifier))))
            for match in self.matches]

        buffer = BytesIO()
        doc = SimpleDocTemplate(
            buffer, pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )
        # COVER PAGE
        self.pdf_elements.extend(
            api.NotificationApi.get_cover_page(
                report_id=report_id,
                recipient=recipient,
            ),
        )

        # MATCH REPORTS
        self.pdf_elements.append(
            Paragraph(
                api.NotificationApi.report_title,
                self.report_title_style,
            )
        )

        # perpetrator info
        self.pdf_elements.append(
            Paragraph(
                "Perpetrator",
                self.section_title_style,
            )
        )
        names = ', '.join(OrderedDict.fromkeys([report.perp_name.strip(
        ) for _, report in matches_with_reports if report.perp_name]))
        if len(names) < 1:
            names = 'None provided'
        perp_info = ('Name(s): ' + names) + '
' + \
            "Matching identifier: " + self.identifier
        self.pdf_elements.append(Paragraph(perp_info, self.body_style))

        # reporter info
        for idx, (match_report, match_report_content) in enumerate(
                matches_with_reports):
            self.pdf_elements.append(
                Paragraph(
                    "Report " + str(idx + 1),
                    self.section_title_style,
                )
            )

            report = match_report.report
            user = report.owner

            if report.submitted_to_school:
                sent_report = report.sentfullreport_set.first()
                report_id = sent_report.get_report_id() if sent_report else "Not found"
                is_submitted = '''
                    Yes
                    Submitted to school on: {0}
                    Submitted report ID: {1}
                '''.format(
                    report.submitted_to_school,
                    report_id,
                )
            else:
                is_submitted = "No"

            overview_body = '''
                Perpetrator name given: {0}
                Reported by: {1}
                Submitted to matching on: {2}
                Record created: {3}
                Full record submitted? {4}
            '''.format(
                match_report_content.perp_name or "None provided",
                self.get_user_identifier(user),
                match_report.added.strftime("%Y-%m-%d %H:%M"),
                report.added.strftime("%Y-%m-%d %H:%M"),
                is_submitted,
            ).replace('\n', '
\n')

            self.pdf_elements.append(Paragraph(overview_body, self.body_style))

            contact_body = '''
Name: {0}
                              Phone: {1}
                              Voicemail preferences: {2}
                              Email: {3}
                              Notes on preferred contact time of day, gender of admin, etc.:'''.format(
                match_report_content.contact_name or "None provided",
                match_report_content.phone,
                match_report_content.voicemail or "None provided",
                match_report_content.email).replace('\n', '
\n')

            self.pdf_elements.append(Paragraph(contact_body, self.body_style))
            self.pdf_elements.append(Paragraph(
                match_report_content.notes or "None provided",
                self.notes_style,
            ))

        doc.build(
            self.pdf_elements,
            onFirstPage=self.get_header_footer(recipient),
            onLaterPages=self.get_header_footer(recipient),
            canvasmaker=NumberedCanvas,
        )
        result = buffer.getvalue()
        buffer.close()
        return result
<|endoftext|>


<commit_before>from slackclient._slackrequest import SlackRequest
from slackclient._channel import Channel
from slackclient._user import User
from slackclient._util import SearchList
from ssl import SSLError

from websocket import create_connection
import json

class Server(object):
    def __init__(self, token, connect=True):
        self.token = token
        self.username = None
        self.domain = None
        self.login_data = None
        self.websocket = None
        self.users = SearchList()
        self.channels = SearchList()
        self.connected = False
        self.pingcounter = 0
        self.api_requester = SlackRequest()

        if connect:
            self.rtm_connect()

    def __eq__(self, compare_str):
        if compare_str == self.domain or compare_str == self.token:
            return True
        else:
            return False

    def __str__(self):
        data = ""
        for key in list(self.__dict__.keys()):
            data += "{} : {}\n".format(key, str(self.__dict__[key])[:40])
        return data

    def __repr__(self):
        return self.__str__()

    def rtm_connect(self, reconnect=False):
        reply = self.api_requester.do(self.token, "rtm.start")
        if reply.code != 200:
            raise SlackConnectionError
        else:
            login_data = json.loads(reply.read().decode('utf-8'))
            if login_data["ok"]:
                self.ws_url = login_data['url']
                if not reconnect:
                    self.parse_slack_login_data(login_data)
                self.connect_slack_websocket(self.ws_url)
            else:
                raise SlackLoginError

    def parse_slack_login_data(self, login_data):
        self.login_data = login_data
        self.domain = self.login_data["team"]["domain"]
        self.username = self.login_data["self"]["name"]
        self.parse_channel_data(login_data["channels"])
        self.parse_channel_data(login_data["groups"])
        self.parse_channel_data(login_data["ims"])
        self.parse_user_data(login_data["users"])

    def connect_slack_websocket(self, ws_url):
        try:
            self.websocket = create_connection(ws_url)
            self.websocket.sock.setblocking(0)
        except:
            raise SlackConnectionError

    def parse_channel_data(self, channel_data):
        for channel in channel_data:
            if "name" not in channel:
                channel["name"] = channel["id"]
            if "members" not in channel:
                channel["members"] = []
            self.attach_channel(channel["name"],
                                channel["id"],
                                channel["members"])

    def parse_user_data(self, user_data):
        for user in user_data:
            if "tz" not in user:
                user["tz"] = "unknown"
            if "real_name" not in user:
                user["real_name"] = user["name"]
            self.attach_user(user["name"], user["id"], user["real_name"], user["tz"])

    def send_to_websocket(self, data):
        """Send (data) directly to the websocket."""
        try:
            data = json.dumps(data)
            self.websocket.send(data)
        except:
            self.rtm_connect(reconnect=True)

    def ping(self):
        return self.send_to_websocket({"type": "ping"})

    def websocket_safe_read(self):
        """ Returns data if available, otherwise ''. Newlines indicate multiple
            messages
        """

        data = ""
        while True:
            try:
                data += "{}\n".format(self.websocket.recv())
            except SSLError as e:
                if e.errno == 2:
                    # errno 2 occurs when trying to read or write data, but more
                    # data needs to be received on the underlying TCP transport
                    # before the request can be fulfilled.
                    #
                    # Python 2.7.9+ and Python 3.3+ give this its own exception,
                    # SSLWantReadError
                    return ''
                raise
            return data.rstrip()

    def attach_user(self, name, id, real_name, tz):
        self.users.append(User(self, name, id, real_name, tz))

    def attach_channel(self, name, id, members=[]):
        self.channels.append(Channel(self, name, id, members))

    def join_channel(self, name):
        print(self.api_requester.do(self.token,
                                    "channels.join?name={}".format(name)).read())

    def api_call(self, method, **kwargs):
        reply = self.api_requester.do(self.token, method, kwargs)
        return reply.read()

class SlackConnectionError(Exception):
    pass

class SlackLoginError(Exception):
    pass
<commit_msg>Fix duplicate channels and users issue.<commit_after>from slackclient._slackrequest import SlackRequest
from slackclient._channel import Channel
from slackclient._user import User
from slackclient._util import SearchList
from ssl import SSLError

from websocket import create_connection
import json

class Server(object):
    def __init__(self, token, connect=True):
        self.token = token
        self.username = None
        self.domain = None
        self.login_data = None
        self.websocket = None
        self.users = SearchList()
        self.channels = SearchList()
        self.connected = False
        self.pingcounter = 0
        self.api_requester = SlackRequest()

        if connect:
            self.rtm_connect()

    def __eq__(self, compare_str):
        if compare_str == self.domain or compare_str == self.token:
            return True
        else:
            return False

    def __str__(self):
        data = ""
        for key in list(self.__dict__.keys()):
            data += "{} : {}\n".format(key, str(self.__dict__[key])[:40])
        return data

    def __repr__(self):
        return self.__str__()

    def rtm_connect(self, reconnect=False):
        reply = self.api_requester.do(self.token, "rtm.start")
        if reply.code != 200:
            raise SlackConnectionError
        else:
            login_data = json.loads(reply.read().decode('utf-8'))
            if login_data["ok"]:
                self.ws_url = login_data['url']
                if not reconnect:
                    self.parse_slack_login_data(login_data)
                self.connect_slack_websocket(self.ws_url)
            else:
                raise SlackLoginError

    def parse_slack_login_data(self, login_data):
        self.login_data = login_data
        self.domain = self.login_data["team"]["domain"]
        self.username = self.login_data["self"]["name"]
        self.parse_channel_data(login_data["channels"])
        self.parse_channel_data(login_data["groups"])
        self.parse_channel_data(login_data["ims"])
        self.parse_user_data(login_data["users"])

    def connect_slack_websocket(self, ws_url):
        try:
            self.websocket = create_connection(ws_url)
            self.websocket.sock.setblocking(0)
        except:
            raise SlackConnectionError

    def parse_channel_data(self, channel_data):
        for channel in channel_data:
            if "name" not in channel:
                channel["name"] = channel["id"]
            if "members" not in channel:
                channel["members"] = []
            self.attach_channel(channel["name"],
                                channel["id"],
                                channel["members"])

    def parse_user_data(self, user_data):
        for user in user_data:
            if "tz" not in user:
                user["tz"] = "unknown"
            if "real_name" not in user:
                user["real_name"] = user["name"]
            self.attach_user(user["name"], user["id"], user["real_name"], user["tz"])

    def send_to_websocket(self, data):
        """Send (data) directly to the websocket."""
        try:
            data = json.dumps(data)
            self.websocket.send(data)
        except:
            self.rtm_connect(reconnect=True)

    def ping(self):
        return self.send_to_websocket({"type": "ping"})

    def websocket_safe_read(self):
        """ Returns data if available, otherwise ''. Newlines indicate multiple
            messages
        """

        data = ""
        while True:
            try:
                data += "{}\n".format(self.websocket.recv())
            except SSLError as e:
                if e.errno == 2:
                    # errno 2 occurs when trying to read or write data, but more
                    # data needs to be received on the underlying TCP transport
                    # before the request can be fulfilled.
                    #
                    # Python 2.7.9+ and Python 3.3+ give this its own exception,
                    # SSLWantReadError
                    return ''
                raise
            return data.rstrip()

    def attach_user(self, name, id, real_name, tz):
        if self.users.find(id) is None:
            self.users.append(User(self, name, id, real_name, tz))

    def attach_channel(self, name, id, members=[]):
        if self.channels.find(id) is None:
            self.channels.append(Channel(self, name, id, members))

    def join_channel(self, name):
        print(self.api_requester.do(self.token,
                                    "channels.join?name={}".format(name)).read())

    def api_call(self, method, **kwargs):
        reply = self.api_requester.do(self.token, method, kwargs)
        return reply.read()

class SlackConnectionError(Exception):
    pass

class SlackLoginError(Exception):
    pass
<|endoftext|>


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

from nose.tools import raises, assert_equal
from ..apis.api import requestURL
import socket

class TestRequestURL:
    @classmethod
    def setup_class(cls):
        cls.oldsocket = socket.socket

    @staticmethod
    def guard(*args, **kwargs):
        raise Exception("Internet access intentionally blocked.")

    def test_empty_json(self):
        branch = "devel"
        url = ("https://raw.githubusercontent.com/Sebelino/SL-CLI/{}/slcli/test"
               "/assets/json/empty.json").format(branch)
        returned = requestURL(url)
        expected = {}
        assert_equal(returned, expected)

    def test_nontrivial_json(self):
        assert False

    def test_non_json(self):
        assert False

    @raises(Exception)
    def test_closed_socket(self):
        socket.socket = self.guard
        requestURL(self.url)

    def teardown(self):
        """ Restore internet connection in case it was blocked """
        socket.socket = self.oldsocket
<commit_msg>Added successful tests for JSON and non-JSON.<commit_after>#!/usr/bin/env python

from nose.tools import raises, assert_equal
from ..apis.api import requestURL
import socket

class TestRequestURL:
    @classmethod
    def setup_class(cls):
        cls.oldsocket = socket.socket

    @staticmethod
    def guard(*args, **kwargs):
        raise Exception("Internet access intentionally blocked.")

    def request_url(self, filename, expected):
        baseurl = ("https://raw.githubusercontent.com/Sebelino/SL-CLI/devel/slc"
                   "li/test/assets/json/")
        url = baseurl+filename
        returned = requestURL(url)
        assert_equal(returned, expected)

    def test_empty_json(self):
        self.request_url("empty.json", dict())

    def test_nontrivial_json(self):
        expected = {
            "glossary": {
                "title": "example glossary",
                "GlossDiv": {
                    "title": "S",
                    "GlossList": {
                        "GlossEntry": {
                            "ID": "SGML",
                            "SortAs": "SGML",
                            "GlossTerm": "Standard Generalized Markup Language",
                            "Acronym": "SGML",
                            "Abbrev": "ISO 8879:1986",
                            "GlossDef": {
                                "para": ("A meta-markup language, used to creat"
                                         "e markup languages such as DocBook."),
                                "GlossSeeAlso": ["GML", "XML"]
                            },
                            "GlossSee": "markup"
                        }
                    }
                }
            }
        }
        self.request_url("nontrivial.json", expected)

    @raises(ValueError)
    def test_non_json(self):
        self.request_url("nonjson.txt", dict())

    @raises(Exception)
    def test_closed_socket(self):
        socket.socket = self.guard
        requestURL(self.url)

    def teardown(self):
        """ Restore internet connection in case it was blocked """
        socket.socket = self.oldsocket
<|endoftext|>


<commit_before>from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk, BulkIndexError
from elasticsearch.client import IndicesClient
import logging

logger = logging.getLogger('sm-engine')

COLUMNS = ["db_name", "ds_id", "ds_name", "sf", "sf_adduct", "comp_names", "comp_ids", "chaos", "image_corr",
           "pattern_match", "msm",
           "adduct", "job_id", "sf_id", "peaks", "db_id", "fdr", "mz", "ds_meta"]

ANNOTATIONS_SEL = '''
SELECT
    sf_db.name AS db_name,
    ds.id as ds_id,
    ds.name AS ds_name,
    f.sf,
    CONCAT(f.sf, m.adduct) as sf_adduct,
    f.names AS comp_names,
    f.subst_ids AS comp_ids,
    COALESCE(((m.stats -> 'chaos'::text)::text)::real, 0::real) AS chaos,
    COALESCE(((m.stats -> 'spatial'::text)::text)::real, 0::real) AS image_corr,
    COALESCE(((m.stats -> 'spectral'::text)::text)::real, 0::real) AS pattern_match,
    COALESCE(m.msm, 0::real) AS msm,
    m.adduct,
    j.id AS job_id,
    f.id AS sf_id,
    m.peaks_n AS peaks,
    sf_db.id AS db_id,
    m.fdr as pass_fdr,
    tp.centr_mzs[1] AS mz,
    ds.metadata as ds_meta
FROM iso_image_metrics m
JOIN formula_db sf_db ON sf_db.id = m.db_id
JOIN sum_formula f ON m.db_id = f.db_id AND f.id = m.sf_id
JOIN job j ON j.id = m.job_id
JOIN dataset ds ON ds.id = j.ds_id
JOIN theor_peaks tp ON tp.sf = f.sf AND tp.adduct = m.adduct
	AND tp.sigma::real = (ds.config->'isotope_generation'->>'isocalc_sigma')::real
	AND tp.charge = (CASE WHEN ds.config->'isotope_generation'->'charge'->>'polarity' = '+' THEN 1 ELSE -1 END)
	AND tp.pts_per_mz = (ds.config->'isotope_generation'->>'isocalc_pts_per_mz')::int
WHERE ds.id = %s
ORDER BY COALESCE(m.msm, 0::real) DESC
'''

class ESExporter:
    def __init__(self, sm_config):
        self.es = Elasticsearch(hosts=[{"host": sm_config['elasticsearch']['host']}])
        self.ind_client = IndicesClient(self.es)
        self.index = sm_config['elasticsearch']['index']

    def _index(self, annotations):
        to_index = []
        for r in annotations:
            d = dict(zip(COLUMNS, r))
            d['comp_names'] = u'|'.join(d['comp_names']).replace(u'"', u'')
            d['comp_ids'] = u'|'.join(d['comp_ids'])
            d['mz'] = '{:010.4f}'.format(d['mz']) if d['mz'] else ''

            to_index.append({
                '_index': self.index,
                '_type': 'annotation',
                '_id': '{}_{}_{}_{}'.format(d['ds_id'], d['db_name'], d['sf'], d['adduct']),
                '_source': d
            })

        bulk(self.es, actions=to_index, timeout='60s')

    def index_ds(self, db, ds_id):
        annotations = db.select(ANNOTATIONS_SEL, ds_id)

        logger.info('Deleting {} documents from the index: {}'.format(len(annotations), ds_id))
        self.delete_ds(ds_id)

        logger.info('Indexing {} documents: {}'.format(len(annotations), ds_id))
        self._index(annotations)

    def delete_ds(self, ds_id):
        body = {
            "query": {
                "constant_score": {
                    "filter": {
                        "bool": {
                            "must": [
                                {"term": {"ds_id": ds_id}}
                            ]
                        }
                    }
                }
            }
        }
        res = self.es.search(index=self.index, body=body, _source=False, size=10**9)['hits']['hits']
        to_del = [{'_op_type': 'delete', '_index': 'sm', '_type': 'annotation', '_id': d['_id']} for d in res]

        del_n = 0
        try:
            del_n, _ = bulk(self.es, to_del, timeout='60s')
        except BulkIndexError as e:
            logger.warning('{} - {}'.format(e.args[0], e.args[1][1]))
        return del_n

    def create_index(self):
        body = {
            "settings": {
                "index": {
                    "number_of_shards": 1,
                    "number_of_replicas": 0,
                    "max_result_window": 2147483647,
                    "analysis": {
                        "analyzer": {
                            "analyzer_keyword": {
                                "tokenizer": "keyword",
                                "filter": "lowercase"
                            }
                        }
                    }
                }
            },
            "mappings": {
                "annotation": {
                    # "dynamic_templates": [{
                    #     "notanalyzed": {
                    #         "match": "*",
                    #         "match_mapping_type": "string",
                    #         "mapping": {
                    #             "type": "string",
                    #             "index": "not_analyzed"
                    #         }
                    #     }
                    # }],
                    "properties": {
                        "db_name": {"type": "string", "index": "not_analyzed"},
                        "ds_id": {"type": "string", "index": "not_analyzed"},
                        "ds_name": {"type": "string", "index": "not_analyzed"},
                        "sf": {"type": "string", "index": "not_analyzed"},
                        "sf_adduct": {"type": "string", "index": "not_analyzed"},
                        "comp_names": {
                            "type": "string",
                            "analyzer": "analyzer_keyword",
                        },
                        "comp_ids": {"type": "string", "index": "not_analyzed"},
                        "chaos": {"type": "float", "index": "not_analyzed"},
                        "image_corr": {"type": "float", "index": "not_analyzed"},
                        "pattern_match": {"type": "float", "index": "not_analyzed"},
                        "msm": {"type": "float", "index": "not_analyzed"},
                        "adduct": {"type": "string", "index": "not_analyzed"},
                        "fdr": {"type": "float", "index": "not_analyzed"},
                        "mz": {"type": "string", "index": "not_analyzed"},
                        # dataset metadata
                        "ds_meta": {
                            "properties": {
                                "Submitted_By": {
                                    "properties": {
                                        "Submitter": {
                                            "properties": {
                                                "Email": {"type": "string", "index": "not_analyzed"}
                                            }
                                        },
                                        "Principal_Investigator": {
                                            "properties": {
                                                "Email": {"type": "string", "index": "not_analyzed"}
                                            }
                                        },
                                        "Institution": {"type": "string", "index": "not_analyzed"}
                                    }
                                },
                                "Sample_Preparation": {
                                    "properties": {
                                        "MALDI_Matrix": {"type": "string", "index": "not_analyzed"},
                                        "MALDI_Matrix_Application": {"type": "string", "index": "not_analyzed"}
                                    }
                                },
                                "Sample_Information": {
                                    "properties": {
                                        "Organism": {"type": "string", "index": "not_analyzed"},
                                        "Condition": {"type": "string", "index": "not_analyzed"}
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        if not self.ind_client.exists(self.index):
            out = self.ind_client.create(index=self.index, body=body)
            logger.info('Index {} created\n{}'.format(self.index, out))
        else:
            logger.info('Index {} already exists'.format(self.index))

    def delete_index(self):
        if self.ind_client.exists(self.index):
            out = self.ind_client.delete(self.index)
            logger.info('Index {} deleted\n{}'.format(self.index, out))
<commit_msg>engine/es_export.ESExporter bug fix: trim comp_names, comp_id field values to avoid ES max_bytes_length_exceeded_exception<commit_after>from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk, BulkIndexError
from elasticsearch.client import IndicesClient
import logging

logger = logging.getLogger('sm-engine')

COLUMNS = ["db_name", "ds_id", "ds_name", "sf", "sf_adduct", "comp_names", "comp_ids", "chaos", "image_corr",
           "pattern_match", "msm",
           "adduct", "job_id", "sf_id", "peaks", "db_id", "fdr", "mz", "ds_meta"]

ANNOTATIONS_SEL = '''
SELECT
    sf_db.name AS db_name,
    ds.id as ds_id,
    ds.name AS ds_name,
    f.sf,
    CONCAT(f.sf, m.adduct) as sf_adduct,
    f.names AS comp_names,
    f.subst_ids AS comp_ids,
    COALESCE(((m.stats -> 'chaos'::text)::text)::real, 0::real) AS chaos,
    COALESCE(((m.stats -> 'spatial'::text)::text)::real, 0::real) AS image_corr,
    COALESCE(((m.stats -> 'spectral'::text)::text)::real, 0::real) AS pattern_match,
    COALESCE(m.msm, 0::real) AS msm,
    m.adduct,
    j.id AS job_id,
    f.id AS sf_id,
    m.peaks_n AS peaks,
    sf_db.id AS db_id,
    m.fdr as pass_fdr,
    tp.centr_mzs[1] AS mz,
    ds.metadata as ds_meta
FROM iso_image_metrics m
JOIN formula_db sf_db ON sf_db.id = m.db_id
JOIN sum_formula f ON m.db_id = f.db_id AND f.id = m.sf_id
JOIN job j ON j.id = m.job_id
JOIN dataset ds ON ds.id = j.ds_id
JOIN theor_peaks tp ON tp.sf = f.sf AND tp.adduct = m.adduct
	AND tp.sigma::real = (ds.config->'isotope_generation'->>'isocalc_sigma')::real
	AND tp.charge = (CASE WHEN ds.config->'isotope_generation'->'charge'->>'polarity' = '+' THEN 1 ELSE -1 END)
	AND tp.pts_per_mz = (ds.config->'isotope_generation'->>'isocalc_pts_per_mz')::int
WHERE ds.id = %s
ORDER BY COALESCE(m.msm, 0::real) DESC
'''

class ESExporter:
    def __init__(self, sm_config):
        self.es = Elasticsearch(hosts=[{"host": sm_config['elasticsearch']['host']}])
        self.ind_client = IndicesClient(self.es)
        self.index = sm_config['elasticsearch']['index']

    def _index(self, annotations):
        to_index = []
        for r in annotations:
            d = dict(zip(COLUMNS, r))
            # trimming is needed for avoiding ES max_bytes_length_exceeded_exception
            d['comp_names'] = u'|'.join(d['comp_names']).replace(u'"', u'')[:32766]
            d['comp_ids'] = u'|'.join(d['comp_ids'])[:32766]
            d['mz'] = '{:010.4f}'.format(d['mz']) if d['mz'] else ''

            to_index.append({
                '_index': self.index,
                '_type': 'annotation',
                '_id': '{}_{}_{}_{}'.format(d['ds_id'], d['db_name'], d['sf'], d['adduct']),
                '_source': d
            })

        bulk(self.es, actions=to_index, timeout='60s')

    def index_ds(self, db, ds_id):
        annotations = db.select(ANNOTATIONS_SEL, ds_id)

        logger.info('Deleting {} documents from the index: {}'.format(len(annotations), ds_id))
        self.delete_ds(ds_id)

        logger.info('Indexing {} documents: {}'.format(len(annotations), ds_id))
        self._index(annotations)

    def delete_ds(self, ds_id):
        body = {
            "query": {
                "constant_score": {
                    "filter": {
                        "bool": {
                            "must": [
                                {"term": {"ds_id": ds_id}}
                            ]
                        }
                    }
                }
            }
        }
        res = self.es.search(index=self.index, body=body, _source=False, size=10**9)['hits']['hits']
        to_del = [{'_op_type': 'delete', '_index': 'sm', '_type': 'annotation', '_id': d['_id']} for d in res]

        del_n = 0
        try:
            del_n, _ = bulk(self.es, to_del, timeout='60s')
        except BulkIndexError as e:
            logger.warning('{} - {}'.format(e.args[0], e.args[1][1]))
        return del_n

    def create_index(self):
        body = {
            "settings": {
                "index": {
                    "number_of_shards": 1,
                    "number_of_replicas": 0,
                    "max_result_window": 2147483647,
                    "analysis": {
                        "analyzer": {
                            "analyzer_keyword": {
                                "tokenizer": "keyword",
                                "filter": "lowercase"
                            }
                        }
                    }
                }
            },
            "mappings": {
                "annotation": {
                    # "dynamic_templates": [{
                    #     "notanalyzed": {
                    #         "match": "*",
                    #         "match_mapping_type": "string",
                    #         "mapping": {
                    #             "type": "string",
                    #             "index": "not_analyzed"
                    #         }
                    #     }
                    # }],
                    "properties": {
                        "db_name": {"type": "string", "index": "not_analyzed"},
                        "ds_id": {"type": "string", "index": "not_analyzed"},
                        "ds_name": {"type": "string", "index": "not_analyzed"},
                        "sf": {"type": "string", "index": "not_analyzed"},
                        "sf_adduct": {"type": "string", "index": "not_analyzed"},
                        "comp_names": {
                            "type": "string",
                            "analyzer": "analyzer_keyword",
                        },
                        "comp_ids": {"type": "string", "index": "not_analyzed"},
                        "chaos": {"type": "float", "index": "not_analyzed"},
                        "image_corr": {"type": "float", "index": "not_analyzed"},
                        "pattern_match": {"type": "float", "index": "not_analyzed"},
                        "msm": {"type": "float", "index": "not_analyzed"},
                        "adduct": {"type": "string", "index": "not_analyzed"},
                        "fdr": {"type": "float", "index": "not_analyzed"},
                        "mz": {"type": "string", "index": "not_analyzed"},
                        # dataset metadata
                        "ds_meta": {
                            "properties": {
                                "Submitted_By": {
                                    "properties": {
                                        "Submitter": {
                                            "properties": {
                                                "Email": {"type": "string", "index": "not_analyzed"}
                                            }
                                        },
                                        "Principal_Investigator": {
                                            "properties": {
                                                "Email": {"type": "string", "index": "not_analyzed"}
                                            }
                                        },
                                        "Institution": {"type": "string", "index": "not_analyzed"}
                                    }
                                },
                                "Sample_Preparation": {
                                    "properties": {
                                        "MALDI_Matrix": {"type": "string", "index": "not_analyzed"},
                                        "MALDI_Matrix_Application": {"type": "string", "index": "not_analyzed"}
                                    }
                                },
                                "Sample_Information": {
                                    "properties": {
                                        "Organism": {"type": "string", "index": "not_analyzed"},
                                        "Condition": {"type": "string", "index": "not_analyzed"}
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        if not self.ind_client.exists(self.index):
            out = self.ind_client.create(index=self.index, body=body)
            logger.info('Index {} created\n{}'.format(self.index, out))
        else:
            logger.info('Index {} already exists'.format(self.index))

    def delete_index(self):
        if self.ind_client.exists(self.index):
            out = self.ind_client.delete(self.index)
            logger.info('Index {} deleted\n{}'.format(self.index, out))
<|endoftext|>


<commit_before>"""
Implementation of the class `CartesianStructuredMesh` and its sub-classes.
"""

import os
import sys
import math
from decimal import Decimal

import numpy
import yaml

class Segment(object):
  """
  Contains information about a segment.
  """

  def __init__(self, data=None, vertices=None):
    """
    Creates the segment vertices.

    Parameters
    ----------
    data: dictionary, optional
      YAML data;
      default: None.
    vertices: 1D array of floats, optional
      vertices along the segment;
      default: None.
    """
    if data:
      self.create_from_yaml_data(data)
    elif numpy.all(vertices):
      self.create_from_vertices(vertices)
    self.nb_divisions = self.vertices.size - 1

  def create_from_vertices(self, vertices):
    """
    Stores vertices.
    """
    self.vertices = vertices

  def create_from_yaml_data(self, data):
    """
    Creates vertices from provided YAML data.

    Parameters
    ----------
    data: dictionary
      YAML data.
    """
    self.start, self.end = data['start'], data['end']
    self.width = data['width']
    ratio = self.get_stretch_ratio(self.width,
                                   stretch_ratio=data['stretchRatio'],
                                   aspect_ratio=data['aspectRatio'],
                                   precision=data['precision'])
    self.stretch_ratio = ratio
    self.vertices = self.get_vertices(reverse=data['reverse'])

  def print_parameters(self):
    """
    Prints the parameters of the segment.
    """
    print('\n[info] printing parameters of segment ...')
    print('\tlimits: [{}, {}]'.format(self.start, self.end))
    print('\treal limits: [{}, {}]'.format(self.vertices[0],
                                           self.vertices[-1]))
    print('\tstretching ratio: {}'.format(self.stretch_ratio))
    print('\taspect ratio: {}'.format(self.aspect_ratio))
    print('\tnumber of divisions: {}\n'.format(self.nb_divisions))

  def get_vertices(self, reverse=False):
    """
    Computes the vertices of the segment.

    Parameters
    ----------
    reverse: boolean, optional
      Set 'True' if you want to reverse the stretching order;
      default: False.

    Returns
    -------
    vertices: 1D array of floats
      The vertices of the segment.
    """
    length = abs(self.end - self.start)
    width = self.width
    ratio = self.stretch_ratio
    if abs(ratio - 1.0) < 1.0E-06:
      # uniform discretization
      self.aspect_ratio = 1.0
      if abs(int(round(length / width)) - length / width) > 1.0E-12:
        print('[error] uniform discretization: '
              'length of the segment should be a multiple of the width '
              'provided')
        sys.exit(0)
      return numpy.arange(self.start, self.end + width / 2.0, width)
    # stretched discretization
    n = int(round(math.log(1.0 - length / width * (1.0 - ratio))
                  / math.log(ratio)))
    widths = numpy.empty(n, dtype=numpy.float64)
    widths[0], widths[1:] = width, ratio
    widths = numpy.cumprod(widths)
    # compute the aspect ratio
    self.aspect_ratio = widths[-1] / widths[0]
    # return the vertices
    if reverse:
      # inverse the stretching ratio and reverse the widths
      self.stretch_ratio = 1.0 / ratio
      return numpy.insert(self.end - numpy.cumsum(widths), 0, self.end)[::-1]
    else:
      return numpy.insert(self.start + numpy.cumsum(widths), 0, self.start)

  def get_stretch_ratio(self, width,
                        stretch_ratio=1.0, aspect_ratio=1.0, precision=6):
    """
    Computes the optimal stretching ratio given a targeted stretching ratio
    or a targeted aspect ratio.

    Parameters
    ----------
    width: float
      Width of the first division.
    stretch_ratio: float, optional
      Targeted stretching ratio;
      default: 1.0.
    aspect_ratio: float, optional
      Targeted aspect ratio between the first and last divisions;
      default: 1.0.
    precision: integer, optional
      Precision of the optimal stretching ratio to compute;
      default: 6.

    Returns
    -------
    ratio: float
      The optimal stretching ratio.
    """
    # if stretching ratio provided
    if abs(stretch_ratio - 1.0) > 1.0E-06:
      return self.compute_optimal_stretch_ratio(width, stretch_ratio,
                                                precision=precision)
    # if aspect ratio provided
    elif abs(aspect_ratio - 1.0) > 1.0E-06:
      ratio = self.compute_stretch_ratio(width, aspect_ratio,
                                         precision=precision)
      return self.compute_optimal_stretch_ratio(width, ratio,
                                                precision=precision)
    # uniform discretization
    else:
      return 1.0

  def compute_stretch_ratio(self, width, aspect_ratio,
                            precision=6):
    """
    Computes the stretching ratio provided the a targeted aspect ratio
    between the first and last divisions of the segment.

    Parameters
    ----------
    width: float
      Width of the first division.
    aspect_ratio: float
      Targeted aspect ratio.
    precision: integer, optional
      Precision of the stretching ratio to compute;
      default: 6.

    Returns
    -------
    ratio: float
      The stretching ratio.
    """
    length = abs(self.end - self.start)
    current_precision = 1
    ratio = 2.0
    while current_precision < precision:
        n = int(round(math.log(1.0 - length / width * (1.0 - ratio))
                      / math.log(ratio)))
        candidate_aspect_ratio = ratio**(n - 1)
        if candidate_aspect_ratio < aspect_ratio:
            ratio += (0.1)**current_precision
            current_precision += 1
        else:
            ratio -= (0.1)**current_precision
    return ratio

  def compute_optimal_stretch_ratio_old(self, width, ratio,
                                        precision=6):
    """
    ***DEPRECATED***
    Computes the optimal stretching ratio provided a targeted one.

    Parameters
    ----------
    width: float
      Width of the first division of the segment.
    ratio: float
      Targeted stretching ratio.
    precision: integer, optional
      Precision of the stretching ratio to compute;
      default: 6.

    Returns
    -------
    ratio: float
      The optimal stretching ratio.
    """
    length = abs(self.end - self.start)
    precision_ratio = abs(Decimal(str(ratio)).as_tuple().exponent)
    while precision_ratio < precision:
      try:
        n = int(round(math.log(1.0 - (1.0 - ratio) * length / width)
                      / math.log(ratio)))
        candidate_length = width * (1.0 - ratio**n) / (1.0 - ratio)
      except:
        candidate_length = 0.0
      if candidate_length < length:
        ratio += (0.1)**precision_ratio
        precision_ratio += 1
      else:
        ratio -= (0.1)**precision_ratio
    return ratio

  def compute_optimal_stretch_ratio(self, width, ratio,
                                    precision=6):
    """
    Computes the optimal stretching ratio provided a targeted one.

    Parameters
    ----------
    width: float
      Width of the first division of the segment.
    ratio: float
      Targeted stretching ratio.
    precision: integer, optional
      Precision of the stretching ratio to compute;
      default: 6.

    Returns
    -------
    ratio: float
      The optimal stretching ratio.
    """
    def geometric_sum(a, r, n):
      """
      Computes the sum of the geometric progression.
      """
      return a * (1.0 - r**n) / (1.0 - r)
    length = abs(self.end - self.start)
    precision_ratio = abs(Decimal(str(ratio)).as_tuple().exponent)
    while precision_ratio < precision:
      n = int(math.log(1.0 - (1.0 - ratio) * length / width)
              / math.log(ratio))
      deviation_inf = abs(length - geometric_sum(width, ratio, n))
      deviation_sup = abs(length - geometric_sum(width, ratio, n + 1))
      precision_ratio += 1
      if deviation_inf < deviation_sup:
        ratio += 0.1**precision_ratio
      else:
        ratio -= 0.1**precision_ratio
    return ratio

  def generate_yaml_info(self):
    """
    Generates a dictionary with segment's information ready for YAML.

    The dictionary contains the end, the number of divisions, and the
    stretching ratio of the segment.
    """
    info = {}
    info['end'] = self.end
    info['cells'] = self.nb_divisions
    info['stretchRatio'] = self.stretch_ratio
    return info

class GridLine(object):
  """
  Contains information about a gridline.
  """

  def __init__(self, data=None, vertices=None, label=None):
    """
    Creates a gridline from provided YAML data or vertices.

    Parameters
    ----------
    data: dictionary, optional
      YAML data about the gridline;
      default: None.
    vertices: 1D array of floats, optional
      Vertices along the gridline;
      default: None.
    label: string, optional
      Label of the direction;
      default: None.
    """
    self.label = label
    self.segments = []
    if data:
      self.create_from_yaml_data(data)
    elif numpy.all(vertices):
      self.create_from_vertices(vertices)
    self.nb_divisions = sum(segment.nb_divisions for segment in self.segments)

  def create_from_vertices(self, vertices):
    """
    Defines the gridline from provided vertices as a single segment.

    Parameters
    ----------
    vertices: 1D array of floats
      The vertices along the gridline.
    """
    self.start, self.end = vertices[0], vertices[-1]
    self.segments.append(Segment(vertices=vertices))

  def create_from_yaml_data(self, data):
    """
    Initializes the gridline parameters and computes its vertices.

    A gridline is defined as a sequence of uniform and/or stretched segments.

    Parameters
    ----------
    data: dictionary
      Parameters of the gridline in a YAML format.
    """
    self.label = data['direction']
    self.start = data['start']
    self.end = data['subDomains'][-1]['end']
    for index, node in enumerate(data['subDomains']):
      # store starting point
      start = (data['start'] if index == 0
               else data['subDomains'][index - 1]['end'])
      data['subDomains'][index]['start'] = start
      # set default parameters if not present
      if 'reverse' not in node.keys():
        data['subDomains'][index]['reverse'] = False
      if 'precision' not in node.keys():
        data['subDomains'][index]['precision'] = 6
      if 'aspectRatio' not in node.keys():
        data['subDomains'][index]['aspectRatio'] = 1.0
      if 'stretchRatio' not in node.keys():
        data['subDomains'][index]['stretchRatio'] = 1.0
      # create a segment
      self.segments.append(Segment(data=node))

  def get_vertices(self, precision=6):
    """
    Gets the vertices removing the repeated values at boundaries
    between consecutive segments.

    Parameters
    ----------
    precision: integer, optional
      Precision used to round the vertices so we can remove repeated values;
      default: 6.

    Returns
    -------
    vertices: 1D array of floats
      The vertices along the gridline.
    """
    return numpy.unique(numpy.concatenate(([numpy.round(segment.vertices,
                                                        precision)
                                            for segment in self.segments])))

  def print_parameters(self):
    """
    Prints parameters of the gridline.
    """
    print('[info] printing parameters of the gridline '
          '{} ...'.format(self.label))
    vertices = self.get_vertices()
    print('\tlimits: [{}, {}]'.format(self.start, self.end))
    print('\treal limits: [{}, {}]'.format(vertices[0], vertices[-1]))
    print('\tnumber of divisions: {}'.format(self.nb_divisions))
    for segment in self.segments:
      segment.print_parameters()

  def generate_yaml_info(self):
    """
    Generates a dictionary with gridline's information ready for YAML.

    The dictionary contains the direction, the start, and segments information
    of the gridline.
    """
    info = {}
    info['direction'] = self.label
    info['start'] = self.start
    info['subDomains'] = []
    for segment in self.segments:
      info['subDomains'].append(segment.generate_yaml_info())
    return info

class CartesianStructuredMesh(object):
  """
  Contains info related to a Cartesian structured mesh (stretched or uniform).
  """

  def __init__(self):
    """
    Instantiates an empty mesh.
    """
    self.gridlines = []

  def get_number_cells(self):
    """
    Gets the number of divisions along each gridline and the total number of
    cells.
    """
    nb_divisions = []
    for gridline in self.gridlines:
      nb_divisions.append(gridline.nb_divisions)
    total = 1
    for nb_division in nb_divisions:
      total *= nb_division
    return total, nb_divisions

  def create(self, data):
    """
    Creates the gridlines.

    Parameters
    ----------
    data: list of dictionaries
      Contains YAML information about the gridlines.
    """
    print('[info] generating Cartesian grid ...')
    for node in data:
      self.gridlines.append(GridLine(data=node))

  def print_parameters(self):
    """Prints parameters of the Cartesian structured mesh."""
    print('[info] printing parameters of the Cartesian structured mesh ...')
    nb_cells, _ = self.get_number_cells()
    print('\tnumber of cells: {}'.format(nb_cells))
    for gridline in self.gridlines:
      gridline.print_parameters()

  def write(self, file_path, precision=6, direction='all'):
    """
    Writes the gridlines into a file.

    The first line of the file contains the number of divisions
    along each gridline.
    Then the vertices along each gridline are written in single column
    starting with the first gridline.

    The last argument allows to output the gridline in one given direction.
    In that case, the stations are listed in a single column.

    Parameters
    ----------
    file_path: string
      Path of the file to write in.
    precision: integer, optional
      Precision at which the vertices will be written;
      default: 6.
    direction: string, optional
      Which gridline to write into file;
      default: 'all' (all directional gridlines).
    """
    print('[info] writing gridlines into {} ...'.format(file_path))
    with open(file_path, 'w') as outfile:
      if direction == 'all':
        _, nb_cells_directions = self.get_number_cells()
        outfile.write('\t'.join(str(nb) for nb in nb_cells_directions) + '\n')
    with open(file_path, 'ab') as outfile:
      if direction in ['x', 'all']:
        numpy.savetxt(outfile,
                      self.gridlines[0].get_vertices(precision=precision))
      if direction in ['y', 'all']:
        numpy.savetxt(outfile,
                      self.gridlines[1].get_vertices(precision=precision))
      if len(self.gridlines) == 3 and direction in ['z', 'all']:
        numpy.savetxt(outfile,
                      self.gridlines[2].get_vertices(precision=precision))

  def read(self, file_path):
    """
    Reads the coordinates from the file.

    Parameters
    ----------
    file_name: string
      Name of file containing grid-node stations along each direction.
    """
    print('[info] reading vertices from {} ...'.format(file_path))
    with open(file_path, 'r') as infile:
      nb_divisions = numpy.array([int(n)
                                  for n in infile.readline().strip().split()])
      vertices = numpy.loadtxt(infile, dtype=numpy.float64)
    vertices = numpy.array(numpy.split(vertices,
                                       numpy.cumsum(nb_divisions[:-1] + 1)))
    labels = ['x', 'y', 'z']
    for index, vertices_gridline in enumerate(vertices):
      self.gridlines.append(GridLine(vertices=vertices_gridline,
                                     label=labels[index]))

  def read_yaml_file(self, file_path):
    """
    Parses the YAML file.

    Parameters
    ----------
    file_path: string
      Path the YAML file.

    Returns
    -------
    data: list of dictionaries
      Parsed YAML information.
    """
    print('[info] reading grid parameters from {} ...'.format(file_path))
    with open(file_path, 'r') as infile:
      return yaml.load(infile)

  def write_yaml_file(self, file_path):
    """
    Writes a YAML readable file with information about the Cartesian structured
    mesh.

    Parameters
    ----------
    file_path: string
      Path of the file to write.
    """
    print('[info] writing grid parameters into {} ...'.format(file_path))
    data = []
    for gridline in self.gridlines:
      data.append(gridline.generate_yaml_info())
    nb_cells, nb_cells_directions = self.get_number_cells()
    with open(file_path, 'w') as outfile:
      outfile.write('# {}\n'.format(os.path.basename(file_path)))
      n_cells_string = 'x'.join(str(nb) for nb in nb_cells_directions)
      outfile.write('# {} = {}\n\n'.format(n_cells_string, nb_cells))
      outfile.write(yaml.dump(data, default_flow_style=False))
<commit_msg>Delete deprecated method<commit_after>"""
Implementation of the class `CartesianStructuredMesh` and its sub-classes.
"""

import os
import sys
import math
from decimal import Decimal

import numpy
import yaml

class Segment(object):
  """
  Contains information about a segment.
  """

  def __init__(self, data=None, vertices=None):
    """
    Creates the segment vertices.

    Parameters
    ----------
    data: dictionary, optional
      YAML data;
      default: None.
    vertices: 1D array of floats, optional
      vertices along the segment;
      default: None.
    """
    if data:
      self.create_from_yaml_data(data)
    elif numpy.all(vertices):
      self.create_from_vertices(vertices)
    self.nb_divisions = self.vertices.size - 1

  def create_from_vertices(self, vertices):
    """
    Stores vertices.
    """
    self.vertices = vertices

  def create_from_yaml_data(self, data):
    """
    Creates vertices from provided YAML data.

    Parameters
    ----------
    data: dictionary
      YAML data.
    """
    self.start, self.end = data['start'], data['end']
    self.width = data['width']
    ratio = self.get_stretch_ratio(self.width,
                                   stretch_ratio=data['stretchRatio'],
                                   aspect_ratio=data['aspectRatio'],
                                   precision=data['precision'])
    self.stretch_ratio = ratio
    self.vertices = self.get_vertices(reverse=data['reverse'])

  def print_parameters(self):
    """
    Prints the parameters of the segment.
    """
    print('\n[info] printing parameters of segment ...')
    print('\tlimits: [{}, {}]'.format(self.start, self.end))
    print('\treal limits: [{}, {}]'.format(self.vertices[0],
                                           self.vertices[-1]))
    print('\tstretching ratio: {}'.format(self.stretch_ratio))
    print('\taspect ratio: {}'.format(self.aspect_ratio))
    print('\tnumber of divisions: {}\n'.format(self.nb_divisions))

  def get_vertices(self, reverse=False):
    """
    Computes the vertices of the segment.

    Parameters
    ----------
    reverse: boolean, optional
      Set 'True' if you want to reverse the stretching order;
      default: False.

    Returns
    -------
    vertices: 1D array of floats
      The vertices of the segment.
    """
    length = abs(self.end - self.start)
    width = self.width
    ratio = self.stretch_ratio
    if abs(ratio - 1.0) < 1.0E-06:
      # uniform discretization
      self.aspect_ratio = 1.0
      if abs(int(round(length / width)) - length / width) > 1.0E-12:
        print('[error] uniform discretization: '
              'length of the segment should be a multiple of the width '
              'provided')
        sys.exit(-1)
      return numpy.arange(self.start, self.end + width / 2.0, width)
    # stretched discretization
    n = int(round(math.log(1.0 - length / width * (1.0 - ratio))
                  / math.log(ratio)))
    widths = numpy.empty(n, dtype=numpy.float64)
    widths[0], widths[1:] = width, ratio
    widths = numpy.cumprod(widths)
    # compute the aspect ratio
    self.aspect_ratio = widths[-1] / widths[0]
    # return the vertices
    if reverse:
      # inverse the stretching ratio and reverse the widths
      self.stretch_ratio = 1.0 / ratio
      return numpy.insert(self.end - numpy.cumsum(widths), 0, self.end)[::-1]
    else:
      return numpy.insert(self.start + numpy.cumsum(widths), 0, self.start)

  def get_stretch_ratio(self, width,
                        stretch_ratio=1.0, aspect_ratio=1.0, precision=6):
    """
    Computes the optimal stretching ratio given a targeted stretching ratio
    or a targeted aspect ratio.

    Parameters
    ----------
    width: float
      Width of the first division.
    stretch_ratio: float, optional
      Targeted stretching ratio;
      default: 1.0.
    aspect_ratio: float, optional
      Targeted aspect ratio between the first and last divisions;
      default: 1.0.
    precision: integer, optional
      Precision of the optimal stretching ratio to compute;
      default: 6.

    Returns
    -------
    ratio: float
      The optimal stretching ratio.
    """
    # if stretching ratio provided
    if abs(stretch_ratio - 1.0) > 1.0E-06:
      return self.compute_optimal_stretch_ratio(width, stretch_ratio,
                                                precision=precision)
    # if aspect ratio provided
    elif abs(aspect_ratio - 1.0) > 1.0E-06:
      ratio = self.compute_stretch_ratio(width, aspect_ratio,
                                         precision=precision)
      return self.compute_optimal_stretch_ratio(width, ratio,
                                                precision=precision)
    # uniform discretization
    else:
      return 1.0

  def compute_stretch_ratio(self, width, aspect_ratio,
                            precision=6):
    """
    Computes the stretching ratio provided the a targeted aspect ratio
    between the first and last divisions of the segment.

    Parameters
    ----------
    width: float
      Width of the first division.
    aspect_ratio: float
      Targeted aspect ratio.
    precision: integer, optional
      Precision of the stretching ratio to compute;
      default: 6.

    Returns
    -------
    ratio: float
      The stretching ratio.
    """
    length = abs(self.end - self.start)
    current_precision = 1
    ratio = 2.0
    while current_precision < precision:
        n = int(round(math.log(1.0 - length / width * (1.0 - ratio))
                      / math.log(ratio)))
        candidate_aspect_ratio = ratio**(n - 1)
        if candidate_aspect_ratio < aspect_ratio:
            ratio += (0.1)**current_precision
            current_precision += 1
        else:
            ratio -= (0.1)**current_precision
    return ratio

  def compute_optimal_stretch_ratio(self, width, ratio,
                                    precision=6):
    """
    Computes the optimal stretching ratio provided a targeted one.

    Parameters
    ----------
    width: float
      Width of the first division of the segment.
    ratio: float
      Targeted stretching ratio.
    precision: integer, optional
      Precision of the stretching ratio to compute;
      default: 6.

    Returns
    -------
    ratio: float
      The optimal stretching ratio.
    """
    def geometric_sum(a, r, n):
      """
      Computes the sum of the geometric progression.
      """
      return a * (1.0 - r**n) / (1.0 - r)
    length = abs(self.end - self.start)
    precision_ratio = abs(Decimal(str(ratio)).as_tuple().exponent)
    while precision_ratio < precision:
      n = int(math.log(1.0 - (1.0 - ratio) * length / width)
              / math.log(ratio))
      deviation_inf = abs(length - geometric_sum(width, ratio, n))
      deviation_sup = abs(length - geometric_sum(width, ratio, n + 1))
      precision_ratio += 1
      if deviation_inf < deviation_sup:
        ratio += 0.1**precision_ratio
      else:
        ratio -= 0.1**precision_ratio
    return ratio

  def generate_yaml_info(self):
    """
    Generates a dictionary with segment's information ready for YAML.

    The dictionary contains the end, the number of divisions, and the
    stretching ratio of the segment.
    """
    info = {}
    info['end'] = self.end
    info['cells'] = self.nb_divisions
    info['stretchRatio'] = self.stretch_ratio
    return info

class GridLine(object):
  """
  Contains information about a gridline.
  """

  def __init__(self, data=None, vertices=None, label=None):
    """
    Creates a gridline from provided YAML data or vertices.

    Parameters
    ----------
    data: dictionary, optional
      YAML data about the gridline;
      default: None.
    vertices: 1D array of floats, optional
      Vertices along the gridline;
      default: None.
    label: string, optional
      Label of the direction;
      default: None.
    """
    self.label = label
    self.segments = []
    if data:
      self.create_from_yaml_data(data)
    elif numpy.all(vertices):
      self.create_from_vertices(vertices)
    self.nb_divisions = sum(segment.nb_divisions for segment in self.segments)

  def create_from_vertices(self, vertices):
    """
    Defines the gridline from provided vertices as a single segment.

    Parameters
    ----------
    vertices: 1D array of floats
      The vertices along the gridline.
    """
    self.start, self.end = vertices[0], vertices[-1]
    self.segments.append(Segment(vertices=vertices))

  def create_from_yaml_data(self, data):
    """
    Initializes the gridline parameters and computes its vertices.

    A gridline is defined as a sequence of uniform and/or stretched segments.

    Parameters
    ----------
    data: dictionary
      Parameters of the gridline in a YAML format.
    """
    self.label = data['direction']
    self.start = data['start']
    self.end = data['subDomains'][-1]['end']
    for index, node in enumerate(data['subDomains']):
      # store starting point
      start = (data['start'] if index == 0
               else data['subDomains'][index - 1]['end'])
      data['subDomains'][index]['start'] = start
      # set default parameters if not present
      if 'reverse' not in node.keys():
        data['subDomains'][index]['reverse'] = False
      if 'precision' not in node.keys():
        data['subDomains'][index]['precision'] = 6
      if 'aspectRatio' not in node.keys():
        data['subDomains'][index]['aspectRatio'] = 1.0
      if 'stretchRatio' not in node.keys():
        data['subDomains'][index]['stretchRatio'] = 1.0
      # create a segment
      self.segments.append(Segment(data=node))

  def get_vertices(self, precision=6):
    """
    Gets the vertices removing the repeated values at boundaries
    between consecutive segments.

    Parameters
    ----------
    precision: integer, optional
      Precision used to round the vertices so we can remove repeated values;
      default: 6.

    Returns
    -------
    vertices: 1D array of floats
      The vertices along the gridline.
    """
    return numpy.unique(numpy.concatenate(([numpy.round(segment.vertices,
                                                        precision)
                                            for segment in self.segments])))

  def print_parameters(self):
    """
    Prints parameters of the gridline.
    """
    print('[info] printing parameters of the gridline '
          '{} ...'.format(self.label))
    vertices = self.get_vertices()
    print('\tlimits: [{}, {}]'.format(self.start, self.end))
    print('\treal limits: [{}, {}]'.format(vertices[0], vertices[-1]))
    print('\tnumber of divisions: {}'.format(self.nb_divisions))
    for segment in self.segments:
      segment.print_parameters()

  def generate_yaml_info(self):
    """
    Generates a dictionary with gridline's information ready for YAML.

    The dictionary contains the direction, the start, and segments information
    of the gridline.
    """
    info = {}
    info['direction'] = self.label
    info['start'] = self.start
    info['subDomains'] = []
    for segment in self.segments:
      info['subDomains'].append(segment.generate_yaml_info())
    return info

class CartesianStructuredMesh(object):
  """
  Contains info related to a Cartesian structured mesh (stretched or uniform).
  """

  def __init__(self):
    """
    Instantiates an empty mesh.
    """
    self.gridlines = []

  def get_number_cells(self):
    """
    Gets the number of divisions along each gridline and the total number of
    cells.
    """
    nb_divisions = []
    for gridline in self.gridlines:
      nb_divisions.append(gridline.nb_divisions)
    total = 1
    for nb_division in nb_divisions:
      total *= nb_division
    return total, nb_divisions

  def create(self, data):
    """
    Creates the gridlines.

    Parameters
    ----------
    data: list of dictionaries
      Contains YAML information about the gridlines.
    """
    print('[info] generating Cartesian grid ...')
    for node in data:
      self.gridlines.append(GridLine(data=node))

  def print_parameters(self):
    """Prints parameters of the Cartesian structured mesh."""
    print('[info] printing parameters of the Cartesian structured mesh ...')
    nb_cells, _ = self.get_number_cells()
    print('\tnumber of cells: {}'.format(nb_cells))
    for gridline in self.gridlines:
      gridline.print_parameters()

  def write(self, file_path, precision=6, direction='all'):
    """
    Writes the gridlines into a file.

    The first line of the file contains the number of divisions
    along each gridline.
    Then the vertices along each gridline are written in single column
    starting with the first gridline.

    The last argument allows to output the gridline in one given direction.
    In that case, the stations are listed in a single column.

    Parameters
    ----------
    file_path: string
      Path of the file to write in.
    precision: integer, optional
      Precision at which the vertices will be written;
      default: 6.
    direction: string, optional
      Which gridline to write into file;
      default: 'all' (all directional gridlines).
    """
    print('[info] writing gridlines into {} ...'.format(file_path))
    with open(file_path, 'w') as outfile:
      if direction == 'all':
        _, nb_cells_directions = self.get_number_cells()
        outfile.write('\t'.join(str(nb) for nb in nb_cells_directions) + '\n')
    with open(file_path, 'ab') as outfile:
      if direction in ['x', 'all']:
        numpy.savetxt(outfile,
                      self.gridlines[0].get_vertices(precision=precision))
      if direction in ['y', 'all']:
        numpy.savetxt(outfile,
                      self.gridlines[1].get_vertices(precision=precision))
      if len(self.gridlines) == 3 and direction in ['z', 'all']:
        numpy.savetxt(outfile,
                      self.gridlines[2].get_vertices(precision=precision))

  def read(self, file_path):
    """
    Reads the coordinates from the file.

    Parameters
    ----------
    file_name: string
      Name of file containing grid-node stations along each direction.
    """
    print('[info] reading vertices from {} ...'.format(file_path))
    with open(file_path, 'r') as infile:
      nb_divisions = numpy.array([int(n)
                                  for n in infile.readline().strip().split()])
      vertices = numpy.loadtxt(infile, dtype=numpy.float64)
    vertices = numpy.array(numpy.split(vertices,
                                       numpy.cumsum(nb_divisions[:-1] + 1)))
    labels = ['x', 'y', 'z']
    for index, vertices_gridline in enumerate(vertices):
      self.gridlines.append(GridLine(vertices=vertices_gridline,
                                     label=labels[index]))

  def read_yaml_file(self, file_path):
    """
    Parses the YAML file.

    Parameters
    ----------
    file_path: string
      Path the YAML file.

    Returns
    -------
    data: list of dictionaries
      Parsed YAML information.
    """
    print('[info] reading grid parameters from {} ...'.format(file_path))
    with open(file_path, 'r') as infile:
      return yaml.load(infile)

  def write_yaml_file(self, file_path):
    """
    Writes a YAML readable file with information about the Cartesian structured
    mesh.

    Parameters
    ----------
    file_path: string
      Path of the file to write.
    """
    print('[info] writing grid parameters into {} ...'.format(file_path))
    data = []
    for gridline in self.gridlines:
      data.append(gridline.generate_yaml_info())
    nb_cells, nb_cells_directions = self.get_number_cells()
    with open(file_path, 'w') as outfile:
      outfile.write('# {}\n'.format(os.path.basename(file_path)))
      n_cells_string = 'x'.join(str(nb) for nb in nb_cells_directions)
      outfile.write('# {} = {}\n\n'.format(n_cells_string, nb_cells))
      outfile.write(yaml.dump(data, default_flow_style=False))
<|endoftext|>