<commit_before>from django.core import serializers
import copy
import json
from pymongo import MongoClient
from django.db import models
from django.db.models.signals import post_save, pre_delete
from django.core.files.storage import FileSystemStorage

import registry.groups.models
from registry.utils import get_working_groups, get_registries

from rdrf.models import Registry

import logging
logger = logging.getLogger('patient')

from registry.utils import stripspaces
from django.conf import settings # for APP_NAME

file_system = FileSystemStorage(location=settings.MEDIA_ROOT, base_url=settings.MEDIA_URL)

class State(models.Model):
    short_name = models.CharField(max_length=3, primary_key=True)
    name = models.CharField(max_length=30)

    class Meta:
        ordering = ["name"]

    def __unicode__(self):
        return self.name

class Doctor(models.Model):
    # TODO: Is it possible for one doctor to work with multiple working groups?
    family_name = models.CharField(max_length=100, db_index=True)
    given_names = models.CharField(max_length=100, db_index=True)
    surgery_name = models.CharField(max_length=100, blank=True)
    speciality = models.CharField(max_length=100)
    address = models.TextField()
    suburb = models.CharField(max_length=50, verbose_name="Suburb/Town")
    state = models.ForeignKey(State, verbose_name="State/Province/Territory")
    phone = models.CharField(max_length=30, blank=True, null=True)
    email = models.EmailField(blank=True, null=True)

    class Meta:
        ordering = ['family_name']

    def __unicode__(self):
        return "%s %s" % (self.family_name.upper(), self.given_names)

class NextOfKinRelationship(models.Model):
    relationship = models.CharField(max_length=100, verbose_name="Relationship")

    class Meta:
        verbose_name = 'Next of Kin Relationship'

    def __unicode__(self):
        return self.relationship

class PatientManager(models.Manager):
    def get_by_registry(self, registry):
        return self.model.objects.filter(rdrf_registry__in=registry)

    def get_by_working_group(self, user):
        return self.model.objects.filter(working_group__in=get_working_groups(user))

    def get_filtered(self, user):
        return self.model.objects.filter(rdrf_registry__id__in=get_registries(user)).filter(working_group__in=get_working_groups(user)).distinct()

    def get_filtered_unallocated(self, user):
        return self.model.objects.filter(working_group__in=get_working_groups(user)).exclude(rdrf_registry__isnull=False)

class Patient(models.Model):
    _MONGO_PATIENT_DATABASE = 'patients'
    _MONGO_PATIENT_COLLECTION = 'patient'

    if settings.INSTALL_NAME == 'dm1':   # Trac #16 item 9
        SEX_CHOICES = ( ("M", "Male"), ("F", "Female") )
    else:
        SEX_CHOICES = ( ("M", "Male"), ("F", "Female"), ("X", "Other/Intersex") )

    objects = PatientManager()
    rdrf_registry = models.ManyToManyField(Registry)
    working_group = models.ForeignKey(registry.groups.models.WorkingGroup, null=False, blank=False, verbose_name="Centre")
    consent = models.BooleanField(null=False, blank=False, help_text="The patient consents to be part of the registry and have data retained and shared in accordance with the information provided to them.", verbose_name="consent given")
    consent_clinical_trials = models.BooleanField(null=False, blank=False, help_text="The patient consents to be contacted about clinical trials or other studies related to their condition.", verbose_name="consent to allow clinical trials given", default=False)
    consent_sent_information = models.BooleanField(null=False, blank=False, help_text="The patient consents to be sent information on their condition.", verbose_name="consent to be sent information given", default=False)
    family_name = models.CharField(max_length=100, db_index=True)
    given_names = models.CharField(max_length=100, db_index=True)
    umrn = models.CharField(max_length=50, null=True, blank=True, db_index=True, verbose_name="Hospital/Clinic ID")
    date_of_birth = models.DateField()
    place_of_birth = models.CharField(max_length=100, null=True, blank=True, verbose_name="Place of Birth")
    date_of_migration = models.DateField(help_text="If migrated", blank=True, null=True)
    sex = models.CharField(max_length=1, choices=SEX_CHOICES)
    home_phone = models.CharField(max_length=30, blank=True, null=True)
    mobile_phone = models.CharField(max_length=30, blank=True, null=True)
    work_phone = models.CharField(max_length=30, blank=True, null=True)
    email = models.EmailField(blank=True, null=True)
    next_of_kin_family_name = models.CharField(max_length=100, blank=True, null=True, verbose_name="family name")
    next_of_kin_given_names = models.CharField(max_length=100, blank=True, null=True, verbose_name="given names")
    next_of_kin_relationship = models.ForeignKey(NextOfKinRelationship, verbose_name="Relationship", blank=True, null=True)
    next_of_kin_address = models.TextField(blank=True, null=True, verbose_name="Address")
    next_of_kin_suburb = models.CharField(max_length=50, blank=True, null=True, verbose_name="Suburb/Town")
    next_of_kin_state = models.ForeignKey(State, verbose_name="State/Province/Territory", related_name="next_of_kin_set", blank=True, null=True)
    next_of_kin_postcode = models.IntegerField(verbose_name="Postcode", blank=True, null=True)
    next_of_kin_home_phone = models.CharField(max_length=30, blank=True, null=True, verbose_name="home phone")
    next_of_kin_mobile_phone = models.CharField(max_length=30, blank=True, null=True, verbose_name="mobile phone")
    next_of_kin_work_phone = models.CharField(max_length=30, blank=True, null=True, verbose_name="work phone")
    next_of_kin_email = models.EmailField(blank=True, null=True, verbose_name="email")
    next_of_kin_parent_place_of_birth = models.CharField(max_length=100, verbose_name="Place of birth of parents", blank=True, null=True)
    doctors = models.ManyToManyField(Doctor, through="PatientDoctor")
    active = models.BooleanField(default=True, help_text="Ticked if active in the registry, ie not a deleted record, or deceased patient.")
    inactive_reason = models.TextField(blank=True, null=True, verbose_name="Reason", help_text="Please provide reason for deactivating the patient")

    class Meta:
        ordering = ["family_name", "given_names", "date_of_birth"]

    def __unicode__(self):
        if self.active:
            return "%s %s" % (self.family_name, self.given_names)
        else:
            return "%s %s (Archived)" % (self.family_name, self.given_names)

    def save(self, *args, **kwargs):
        if hasattr(self, 'family_name'):
            self.family_name = stripspaces(self.family_name).upper()

        if hasattr(self, 'given_names'):
            self.given_names = stripspaces(self.given_names)

        if not self.pk:
            self.active = True

        super(Patient, self).save(*args, **kwargs)
        regs = self._save_patient_mongo()

    def _save_patient_mongo(self):
        client = MongoClient()
        patient_db = client[self._MONGO_PATIENT_DATABASE]
        patient_coll = patient_db[self._MONGO_PATIENT_COLLECTION]

        json_str  = serializers.serialize("json", [self,])
        json_obj = json.loads(json_str)

        mongo_doc = patient_coll.find_one({'django_id': json_obj[0]['pk']})

        if(mongo_doc):
            self._update_mongo_obj(mongo_doc, json_obj[0]['fields'])
            patient_coll.save(mongo_doc)
        else:
            json_obj[0]['fields']['django_id'] = json_obj[0]['pk']
            patient_coll.save(json_obj[0]['fields'])

    def _update_mongo_obj(self, mongo_doc, patient_model):
        for key, value in mongo_doc.iteritems():
            if key not in ['django_id', '_id']:
                mongo_doc[key] = patient_model[key]

    def delete(self, *args, **kwargs):
        """
        If a user deletes a patient it's active flag will be true, so we should set it to false.
        If a superuser deletes a patient it's active flag is false, so we should delete the object.
        """
        if self.active:
            logger.debug("Archiving patient record.")
            self.active = False
            self.save()
        else:
            logger.debug("Deleting patient record.")
            super(Patient, self).delete(*args, **kwargs)

    def get_reg_list(self):
        return ', '.join([r.name for r in self.rdrf_registry.all()])
    get_reg_list.short_description = 'Registry'

    def as_json(self):
        return dict(
            obj_id=self.id,
            given_names=self.given_names,
            family_name=self.family_name,
            working_group=self.working_group.name,
            date_of_birth=str(self.date_of_birth)
            )

class AddressType(models.Model):
    type = models.CharField(max_length=100)
    description = models.TextField(null=True, blank=True)

    def __unicode__(self):
        return "%s" % (self.type)

class PatientAddress(models.Model):
    patient = models.ForeignKey(Patient)
    address_type = models.ForeignKey(AddressType, default=1)
    address = models.TextField()
    suburb = models.CharField(max_length=50, verbose_name="Suburb/Town")
    state = models.CharField(max_length=20, verbose_name="State/Province/Territory")
    postcode = models.IntegerField()
    country = models.CharField(max_length=20)

    class Meta:
        verbose_name_plural = "Patient Addresses"

class PatientConsent(models.Model):
    patient = models.ForeignKey(Patient)
    form = models.FileField(upload_to='consents', storage=file_system, verbose_name="Consent form", blank=True, null=True)

class PatientDoctor(models.Model):
    patient = models.ForeignKey(Patient)
    doctor = models.ForeignKey(Doctor)
    relationship = models.CharField(max_length=50)

    class Meta:
        verbose_name = "medical professionals for patient"
        verbose_name_plural = "medical professionals for patient"
<commit_msg>Commented out patient save to mongo -<commit_after>from django.core import serializers
import copy
import json
from pymongo import MongoClient
from django.db import models
from django.db.models.signals import post_save, pre_delete
from django.core.files.storage import FileSystemStorage

import registry.groups.models
from registry.utils import get_working_groups, get_registries

from rdrf.models import Registry

import logging
logger = logging.getLogger('patient')

from registry.utils import stripspaces
from django.conf import settings # for APP_NAME

file_system = FileSystemStorage(location=settings.MEDIA_ROOT, base_url=settings.MEDIA_URL)

class State(models.Model):
    short_name = models.CharField(max_length=3, primary_key=True)
    name = models.CharField(max_length=30)

    class Meta:
        ordering = ["name"]

    def __unicode__(self):
        return self.name

class Doctor(models.Model):
    # TODO: Is it possible for one doctor to work with multiple working groups?
    family_name = models.CharField(max_length=100, db_index=True)
    given_names = models.CharField(max_length=100, db_index=True)
    surgery_name = models.CharField(max_length=100, blank=True)
    speciality = models.CharField(max_length=100)
    address = models.TextField()
    suburb = models.CharField(max_length=50, verbose_name="Suburb/Town")
    state = models.ForeignKey(State, verbose_name="State/Province/Territory")
    phone = models.CharField(max_length=30, blank=True, null=True)
    email = models.EmailField(blank=True, null=True)

    class Meta:
        ordering = ['family_name']

    def __unicode__(self):
        return "%s %s" % (self.family_name.upper(), self.given_names)

class NextOfKinRelationship(models.Model):
    relationship = models.CharField(max_length=100, verbose_name="Relationship")

    class Meta:
        verbose_name = 'Next of Kin Relationship'

    def __unicode__(self):
        return self.relationship

class PatientManager(models.Manager):
    def get_by_registry(self, registry):
        return self.model.objects.filter(rdrf_registry__in=registry)

    def get_by_working_group(self, user):
        return self.model.objects.filter(working_group__in=get_working_groups(user))

    def get_filtered(self, user):
        return self.model.objects.filter(rdrf_registry__id__in=get_registries(user)).filter(working_group__in=get_working_groups(user)).distinct()

    def get_filtered_unallocated(self, user):
        return self.model.objects.filter(working_group__in=get_working_groups(user)).exclude(rdrf_registry__isnull=False)

class Patient(models.Model):
    _MONGO_PATIENT_DATABASE = 'patients'
    _MONGO_PATIENT_COLLECTION = 'patient'

    if settings.INSTALL_NAME == 'dm1':   # Trac #16 item 9
        SEX_CHOICES = ( ("M", "Male"), ("F", "Female") )
    else:
        SEX_CHOICES = ( ("M", "Male"), ("F", "Female"), ("X", "Other/Intersex") )

    objects = PatientManager()
    rdrf_registry = models.ManyToManyField(Registry)
    working_group = models.ForeignKey(registry.groups.models.WorkingGroup, null=False, blank=False, verbose_name="Centre")
    consent = models.BooleanField(null=False, blank=False, help_text="The patient consents to be part of the registry and have data retained and shared in accordance with the information provided to them.", verbose_name="consent given")
    consent_clinical_trials = models.BooleanField(null=False, blank=False, help_text="The patient consents to be contacted about clinical trials or other studies related to their condition.", verbose_name="consent to allow clinical trials given", default=False)
    consent_sent_information = models.BooleanField(null=False, blank=False, help_text="The patient consents to be sent information on their condition.", verbose_name="consent to be sent information given", default=False)
    family_name = models.CharField(max_length=100, db_index=True)
    given_names = models.CharField(max_length=100, db_index=True)
    umrn = models.CharField(max_length=50, null=True, blank=True, db_index=True, verbose_name="Hospital/Clinic ID")
    date_of_birth = models.DateField()
    place_of_birth = models.CharField(max_length=100, null=True, blank=True, verbose_name="Place of Birth")
    date_of_migration = models.DateField(help_text="If migrated", blank=True, null=True)
    sex = models.CharField(max_length=1, choices=SEX_CHOICES)
    home_phone = models.CharField(max_length=30, blank=True, null=True)
    mobile_phone = models.CharField(max_length=30, blank=True, null=True)
    work_phone = models.CharField(max_length=30, blank=True, null=True)
    email = models.EmailField(blank=True, null=True)
    next_of_kin_family_name = models.CharField(max_length=100, blank=True, null=True, verbose_name="family name")
    next_of_kin_given_names = models.CharField(max_length=100, blank=True, null=True, verbose_name="given names")
    next_of_kin_relationship = models.ForeignKey(NextOfKinRelationship, verbose_name="Relationship", blank=True, null=True)
    next_of_kin_address = models.TextField(blank=True, null=True, verbose_name="Address")
    next_of_kin_suburb = models.CharField(max_length=50, blank=True, null=True, verbose_name="Suburb/Town")
    next_of_kin_state = models.ForeignKey(State, verbose_name="State/Province/Territory", related_name="next_of_kin_set", blank=True, null=True)
    next_of_kin_postcode = models.IntegerField(verbose_name="Postcode", blank=True, null=True)
    next_of_kin_home_phone = models.CharField(max_length=30, blank=True, null=True, verbose_name="home phone")
    next_of_kin_mobile_phone = models.CharField(max_length=30, blank=True, null=True, verbose_name="mobile phone")
    next_of_kin_work_phone = models.CharField(max_length=30, blank=True, null=True, verbose_name="work phone")
    next_of_kin_email = models.EmailField(blank=True, null=True, verbose_name="email")
    next_of_kin_parent_place_of_birth = models.CharField(max_length=100, verbose_name="Place of birth of parents", blank=True, null=True)
    doctors = models.ManyToManyField(Doctor, through="PatientDoctor")
    active = models.BooleanField(default=True, help_text="Ticked if active in the registry, ie not a deleted record, or deceased patient.")
    inactive_reason = models.TextField(blank=True, null=True, verbose_name="Reason", help_text="Please provide reason for deactivating the patient")

    class Meta:
        ordering = ["family_name", "given_names", "date_of_birth"]

    def __unicode__(self):
        if self.active:
            return "%s %s" % (self.family_name, self.given_names)
        else:
            return "%s %s (Archived)" % (self.family_name, self.given_names)

    def save(self, *args, **kwargs):
        if hasattr(self, 'family_name'):
            self.family_name = stripspaces(self.family_name).upper()

        if hasattr(self, 'given_names'):
            self.given_names = stripspaces(self.given_names)

        if not self.pk:
            self.active = True

        super(Patient, self).save(*args, **kwargs)
        #regs = self._save_patient_mongo()

    def _save_patient_mongo(self):
        client = MongoClient()
        patient_db = client[self._MONGO_PATIENT_DATABASE]
        patient_coll = patient_db[self._MONGO_PATIENT_COLLECTION]

        json_str  = serializers.serialize("json", [self,])
        json_obj = json.loads(json_str)

        mongo_doc = patient_coll.find_one({'django_id': json_obj[0]['pk']})

        if(mongo_doc):
            self._update_mongo_obj(mongo_doc, json_obj[0]['fields'])
            patient_coll.save(mongo_doc)
        else:
            json_obj[0]['fields']['django_id'] = json_obj[0]['pk']
            patient_coll.save(json_obj[0]['fields'])

    def _update_mongo_obj(self, mongo_doc, patient_model):
        for key, value in mongo_doc.iteritems():
            if key not in ['django_id', '_id']:
                mongo_doc[key] = patient_model[key]

    def delete(self, *args, **kwargs):
        """
        If a user deletes a patient it's active flag will be true, so we should set it to false.
        If a superuser deletes a patient it's active flag is false, so we should delete the object.
        """
        if self.active:
            logger.debug("Archiving patient record.")
            self.active = False
            self.save()
        else:
            logger.debug("Deleting patient record.")
            super(Patient, self).delete(*args, **kwargs)

    def get_reg_list(self):
        return ', '.join([r.name for r in self.rdrf_registry.all()])
    get_reg_list.short_description = 'Registry'

    def as_json(self):
        return dict(
            obj_id=self.id,
            given_names=self.given_names,
            family_name=self.family_name,
            working_group=self.working_group.name,
            date_of_birth=str(self.date_of_birth)
            )

class AddressType(models.Model):
    type = models.CharField(max_length=100)
    description = models.TextField(null=True, blank=True)

    def __unicode__(self):
        return "%s" % (self.type)

class PatientAddress(models.Model):
    patient = models.ForeignKey(Patient)
    address_type = models.ForeignKey(AddressType, default=1)
    address = models.TextField()
    suburb = models.CharField(max_length=50, verbose_name="Suburb/Town")
    state = models.CharField(max_length=20, verbose_name="State/Province/Territory")
    postcode = models.IntegerField()
    country = models.CharField(max_length=20)

    class Meta:
        verbose_name_plural = "Patient Addresses"

class PatientConsent(models.Model):
    patient = models.ForeignKey(Patient)
    form = models.FileField(upload_to='consents', storage=file_system, verbose_name="Consent form", blank=True, null=True)

class PatientDoctor(models.Model):
    patient = models.ForeignKey(Patient)
    doctor = models.ForeignKey(Doctor)
    relationship = models.CharField(max_length=50)

    class Meta:
        verbose_name = "medical professionals for patient"
        verbose_name_plural = "medical professionals for patient"
<|endoftext|>


<commit_before>"""Functions related to converting content into dict/JSON structures."""

import itertools
import logging
from urllib.parse import urlparse

import orjson as json
from django.conf import settings
from django.core.files.storage import get_storage_class
from selectolax.parser import HTMLParser

log = logging.getLogger(__name__)

def generate_page_sections(page_title, body, fjson_storage_path):
    """
    Generate section dicts for each section for sphinx.

    In Sphinx sub-sections are nested, so they are children of the outer section,
    and sections with the same level are neighbors.
    We index the content under a section till before the next one.

    We can have pages that have content before the first title or that don't have a title,
    we index that content first under the title of the original page (`page_title`).

    Contents that are likely to be a sphinx domain are deleted,
    since we already index those in another step.
    """

    # Removing all <dl> tags to prevent duplicate indexing with Sphinx Domains.
    nodes_to_be_removed = []
    try:
        # remove all <dl> tags which contains <dt> tags having 'id' attribute
        dt_tags = body.css('dt[id]')
        for tag in dt_tags:
            parent = tag.parent
            if parent.tag == 'dl':
                nodes_to_be_removed.append(parent)
    except Exception:
        log.exception('Error removing <dl> tags from file: %s', fjson_storage_path)

    # remove `Table of Contents` elements
    try:
        nodes_to_be_removed += body.css('.toctree-wrapper') + body.css('.contents.local.topic')
    except Exception:
        log.exception('Error removing "Table of Content" tags from file: %s', fjson_storage_path)

    # removing all nodes in list
    for node in nodes_to_be_removed:
        node.decompose()

    # Index content for pages that don't start with a title.
    content = _get_content_from_tag(body.body.child)
    if content:
        yield {
            'id': '',
            'title': page_title,
            'content': content,
        }

    # Index content from h1 to h6 headers.
    for head_level in range(1, 7):
        tags = body.css(f'.section > h{head_level}')
        for tag in tags:
            title = _parse_title(tag)

            div = tag.parent
            section_id = div.attributes.get('id', '')

            yield {
                'id': section_id,
                'title': title,
                'content': _get_content_from_tag(tag.next),
            }

def _get_content_from_tag(tag):
    """Gets the content from tag till before a new section."""
    contents = []
    next_tag = tag
    while next_tag and not _is_section(next_tag):
        if _is_code_section(next_tag):
            content = _parse_code_section(next_tag)
        else:
            content = parse_content(next_tag.text())

        if content:
            contents.append(content)
        next_tag = next_tag.next
    return ' '.join(contents)

def _is_code_section(tag):
    """
    Check if `tag` is a code section.

    Sphinx and Mkdocs codeblocks have a class named
    ``highlight`` or ``highlight-{language}``.
    """
    for c in tag.attributes.get('class', '').split():
        if c.startswith('highlight'):
            return True
    return False

def _is_section(tag):
    """Check if `tag` is a sphinx section (linkeable header)."""
    return (
        tag.tag == 'div' and
        'section' in tag.attributes.get('class', [])
    )

def _parse_code_section(tag):
    """
    Parse a code section to fetch relevant content only.

    - Removes line numbers.
      Sphinx and Mkdocs may use a table when the code block includes line numbers.
      This table has a td tag with a ``lineos`` class.
      Other implementations put the line number within the code,
      inside span tags with the ``lineno`` class.
    """
    nodes_to_be_removed = itertools.chain(tag.css('.linenos'), tag.css('.lineno'))
    for node in nodes_to_be_removed:
        node.decompose()

    contents = []
    for node in tag.css('pre'):
        # XXX: Don't call to `parse_content`
        # if we decide to show code results more nicely,
        # so the indentation isn't lost.
        content = node.text().strip('\n')
        contents.append(parse_content(content))
    return ' '.join(contents)

def process_file(fjson_storage_path):
    """Read the fjson file from disk and parse it into a structured dict."""
    storage = get_storage_class(settings.RTD_BUILD_MEDIA_STORAGE)()

    log.debug('Processing JSON file for indexing: %s', fjson_storage_path)

    try:
        with storage.open(fjson_storage_path, mode='r') as f:
            file_contents = f.read()
    except IOError:
        log.info('Unable to read file: %s', fjson_storage_path)
        raise

    data = json.loads(file_contents)
    sections = []
    path = ''
    title = ''
    domain_data = {}

    if 'current_page_name' in data:
        path = data['current_page_name']
    else:
        log.info('Unable to index file due to no name %s', fjson_storage_path)

    if 'title' in data:
        title = data['title']
        title = HTMLParser(title).text().replace('¶', '').strip()
    else:
        log.info('Unable to index title for: %s', fjson_storage_path)

    if data.get('body'):
        body = HTMLParser(data['body'])
        body_copy = HTMLParser(data['body'])
        sections = generate_page_sections(
            page_title=title,
            body=body,
            fjson_storage_path=fjson_storage_path,
        )

        # pass a copy of `body` so that the removed
        # nodes in the original don't reflect here.
        domain_data = generate_domains_data(body_copy, fjson_storage_path)
    else:
        log.info('Unable to index content for: %s', fjson_storage_path)

    return {
        'path': path,
        'title': title,
        'sections': list(sections),
        'domain_data': domain_data,
    }

def generate_domains_data(body, fjson_storage_path):
    """
    Generate sphinx domain objects' docstrings.

    Returns a dict with the generated data.
    The returned dict is in the following form::

        {
            "domain-id-1": "docstrings for the domain-id-1",
            "domain-id-2": "docstrings for the domain-id-2",
        }
    """

    domain_data = {}
    dl_tags = body.css('dl')

    for dl_tag in dl_tags:

        dt = dl_tag.css('dt')
        dd = dl_tag.css('dd')

        # len(dt) should be equal to len(dd)
        # because these tags go together.
        for title, desc in zip(dt, dd):
            try:
                id_ = title.attributes.get('id', '')
                if id_:
                    docstrings = _get_text_for_domain_data(desc)
                    domain_data[id_] = docstrings
            except Exception:
                log.exception('Error parsing docstrings for domains in file %s', fjson_storage_path)

    return domain_data

def _get_text_for_domain_data(desc):
    """Returns the text from the dom node."""

    # remove the 'dl', 'dt' and 'dd' tags from it
    # because all the 'dd' and 'dt' tags are inside 'dl'
    # and all 'dl' tags are already captured.
    nodes_to_be_removed = desc.css('dl') + desc.css('dt') + desc.css('dd')
    for node in nodes_to_be_removed:
        node.decompose()

    # remove multiple spaces, new line characters and '¶' symbol.
    docstrings = parse_content(desc.text())
    return docstrings

def parse_content(content):
    """Removes new line characters and ¶."""
    content = content.replace('¶', '').strip()
    content = content.split('\n')

    # Convert all new lines to " "
    content = (text.strip() for text in content)
    content = ' '.join(text for text in content if text)
    return content

def _parse_title(tag):
    """
    Parses a Sphinx title tag.

    - Removes the permalink value
    """
    nodes_to_be_removed = tag.css('a.headerlink')
    for node in nodes_to_be_removed:
        node.decompose()
    return tag.text().strip()

def process_mkdocs_index_file(json_storage_path, page):
    """Reads the json index file and parses it into a structured dict."""
    log.debug('Processing JSON index file: %s', json_storage_path)

    storage = get_storage_class(settings.RTD_BUILD_MEDIA_STORAGE)()
    try:
        with storage.open(json_storage_path, mode='r') as f:
            file_contents = f.read()
    except IOError:
        log.info('Unable to read file: %s', json_storage_path)
        raise

    data = json.loads(file_contents)
    page_data = {}

    for section in data.get('docs', []):
        parsed_path = urlparse(section.get('location', ''))
        fragment = parsed_path.fragment
        path = parsed_path.path

        # Some old versions of mkdocs
        # index the pages as ``/page.html`` insted of ``page.html``.
        path = path.lstrip('/')

        if path == '' or path.endswith('/'):
            path += 'index.html'

        if page != path:
            continue

        title = HTMLParser(section.get('title')).text()
        content = parse_content(
            HTMLParser(section.get('text')).text()
        )

        # If it doesn't have a fragment,
        # it means is the page itself.
        if not fragment:
            page_data.update({
                'path': path,
                'title': title,
                'domain_data': {},
            })
        # Content without a fragment need to be indexed as well,
        # this happens when the page doesn't start with a header,
        # or if it doesn't contain any headers at all.
        page_data.setdefault('sections', []).append({
            'id': fragment,
            'title': title,
            'content': content,
        })

    return page_data
<commit_msg>Be more defensive<commit_after>"""Functions related to converting content into dict/JSON structures."""

import itertools
import logging
from urllib.parse import urlparse

import orjson as json
from django.conf import settings
from django.core.files.storage import get_storage_class
from selectolax.parser import HTMLParser

log = logging.getLogger(__name__)

def generate_page_sections(page_title, body, fjson_storage_path):
    """
    Generate section dicts for each section for sphinx.

    In Sphinx sub-sections are nested, so they are children of the outer section,
    and sections with the same level are neighbors.
    We index the content under a section till before the next one.

    We can have pages that have content before the first title or that don't have a title,
    we index that content first under the title of the original page (`page_title`).

    Contents that are likely to be a sphinx domain are deleted,
    since we already index those in another step.
    """

    # Removing all <dl> tags to prevent duplicate indexing with Sphinx Domains.
    nodes_to_be_removed = []
    try:
        # remove all <dl> tags which contains <dt> tags having 'id' attribute
        dt_tags = body.css('dt[id]')
        for tag in dt_tags:
            parent = tag.parent
            if parent.tag == 'dl':
                nodes_to_be_removed.append(parent)
    except Exception:
        log.exception('Error removing <dl> tags from file: %s', fjson_storage_path)

    # remove `Table of Contents` elements
    try:
        nodes_to_be_removed += body.css('.toctree-wrapper') + body.css('.contents.local.topic')
    except Exception:
        log.exception('Error removing "Table of Content" tags from file: %s', fjson_storage_path)

    # removing all nodes in list
    for node in nodes_to_be_removed:
        node.decompose()

    # Index content for pages that don't start with a title.
    content = _get_content_from_tag(body.body.child)
    if content:
        yield {
            'id': '',
            'title': page_title,
            'content': content,
        }

    # Index content from h1 to h6 headers.
    for head_level in range(1, 7):
        tags = body.css(f'.section > h{head_level}')
        for tag in tags:
            title = _parse_title(tag)

            div = tag.parent
            section_id = div.attributes.get('id', '')

            yield {
                'id': section_id,
                'title': title,
                'content': _get_content_from_tag(tag.next),
            }

def _get_content_from_tag(tag):
    """Gets the content from tag till before a new section."""
    contents = []
    next_tag = tag
    while next_tag and not _is_section(next_tag):
        if _is_code_section(next_tag):
            content = _parse_code_section(next_tag)
        else:
            content = parse_content(next_tag.text())

        if content:
            contents.append(content)
        next_tag = next_tag.next
    return ' '.join(contents)

def _is_code_section(tag):
    """
    Check if `tag` is a code section.

    Sphinx and Mkdocs codeblocks have a class named
    ``highlight`` or ``highlight-{language}``.
    """
    if not tag.css_first('pre'):
        return False

    for c in tag.attributes.get('class', '').split():
        if c.startswith('highlight'):
            return True
    return False

def _is_section(tag):
    """Check if `tag` is a sphinx section (linkeable header)."""
    return (
        tag.tag == 'div' and
        'section' in tag.attributes.get('class', [])
    )

def _parse_code_section(tag):
    """
    Parse a code section to fetch relevant content only.

    - Removes line numbers.
      Sphinx and Mkdocs may use a table when the code block includes line numbers.
      This table has a td tag with a ``lineos`` class.
      Other implementations put the line number within the code,
      inside span tags with the ``lineno`` class.
    """
    nodes_to_be_removed = itertools.chain(tag.css('.linenos'), tag.css('.lineno'))
    for node in nodes_to_be_removed:
        node.decompose()

    contents = []
    for node in tag.css('pre'):
        # XXX: Don't call to `parse_content`
        # if we decide to show code results more nicely,
        # so the indentation isn't lost.
        content = node.text().strip('\n')
        contents.append(parse_content(content))
    return ' '.join(contents)

def process_file(fjson_storage_path):
    """Read the fjson file from disk and parse it into a structured dict."""
    storage = get_storage_class(settings.RTD_BUILD_MEDIA_STORAGE)()

    log.debug('Processing JSON file for indexing: %s', fjson_storage_path)

    try:
        with storage.open(fjson_storage_path, mode='r') as f:
            file_contents = f.read()
    except IOError:
        log.info('Unable to read file: %s', fjson_storage_path)
        raise

    data = json.loads(file_contents)
    sections = []
    path = ''
    title = ''
    domain_data = {}

    if 'current_page_name' in data:
        path = data['current_page_name']
    else:
        log.info('Unable to index file due to no name %s', fjson_storage_path)

    if 'title' in data:
        title = data['title']
        title = HTMLParser(title).text().replace('¶', '').strip()
    else:
        log.info('Unable to index title for: %s', fjson_storage_path)

    if data.get('body'):
        body = HTMLParser(data['body'])
        body_copy = HTMLParser(data['body'])
        sections = generate_page_sections(
            page_title=title,
            body=body,
            fjson_storage_path=fjson_storage_path,
        )

        # pass a copy of `body` so that the removed
        # nodes in the original don't reflect here.
        domain_data = generate_domains_data(body_copy, fjson_storage_path)
    else:
        log.info('Unable to index content for: %s', fjson_storage_path)

    return {
        'path': path,
        'title': title,
        'sections': list(sections),
        'domain_data': domain_data,
    }

def generate_domains_data(body, fjson_storage_path):
    """
    Generate sphinx domain objects' docstrings.

    Returns a dict with the generated data.
    The returned dict is in the following form::

        {
            "domain-id-1": "docstrings for the domain-id-1",
            "domain-id-2": "docstrings for the domain-id-2",
        }
    """

    domain_data = {}
    dl_tags = body.css('dl')

    for dl_tag in dl_tags:

        dt = dl_tag.css('dt')
        dd = dl_tag.css('dd')

        # len(dt) should be equal to len(dd)
        # because these tags go together.
        for title, desc in zip(dt, dd):
            try:
                id_ = title.attributes.get('id', '')
                if id_:
                    docstrings = _get_text_for_domain_data(desc)
                    domain_data[id_] = docstrings
            except Exception:
                log.exception('Error parsing docstrings for domains in file %s', fjson_storage_path)

    return domain_data

def _get_text_for_domain_data(desc):
    """Returns the text from the dom node."""

    # remove the 'dl', 'dt' and 'dd' tags from it
    # because all the 'dd' and 'dt' tags are inside 'dl'
    # and all 'dl' tags are already captured.
    nodes_to_be_removed = desc.css('dl') + desc.css('dt') + desc.css('dd')
    for node in nodes_to_be_removed:
        node.decompose()

    # remove multiple spaces, new line characters and '¶' symbol.
    docstrings = parse_content(desc.text())
    return docstrings

def parse_content(content):
    """Removes new line characters and ¶."""
    content = content.replace('¶', '').strip()
    content = content.split('\n')

    # Convert all new lines to " "
    content = (text.strip() for text in content)
    content = ' '.join(text for text in content if text)
    return content

def _parse_title(tag):
    """
    Parses a Sphinx title tag.

    - Removes the permalink value
    """
    nodes_to_be_removed = tag.css('a.headerlink')
    for node in nodes_to_be_removed:
        node.decompose()
    return tag.text().strip()

def process_mkdocs_index_file(json_storage_path, page):
    """Reads the json index file and parses it into a structured dict."""
    log.debug('Processing JSON index file: %s', json_storage_path)

    storage = get_storage_class(settings.RTD_BUILD_MEDIA_STORAGE)()
    try:
        with storage.open(json_storage_path, mode='r') as f:
            file_contents = f.read()
    except IOError:
        log.info('Unable to read file: %s', json_storage_path)
        raise

    data = json.loads(file_contents)
    page_data = {}

    for section in data.get('docs', []):
        parsed_path = urlparse(section.get('location', ''))
        fragment = parsed_path.fragment
        path = parsed_path.path

        # Some old versions of mkdocs
        # index the pages as ``/page.html`` insted of ``page.html``.
        path = path.lstrip('/')

        if path == '' or path.endswith('/'):
            path += 'index.html'

        if page != path:
            continue

        title = HTMLParser(section.get('title')).text()
        content = parse_content(
            HTMLParser(section.get('text')).text()
        )

        # If it doesn't have a fragment,
        # it means is the page itself.
        if not fragment:
            page_data.update({
                'path': path,
                'title': title,
                'domain_data': {},
            })
        # Content without a fragment need to be indexed as well,
        # this happens when the page doesn't start with a header,
        # or if it doesn't contain any headers at all.
        page_data.setdefault('sections', []).append({
            'id': fragment,
            'title': title,
            'content': content,
        })

    return page_data
<|endoftext|>


<commit_before>"pymel logging functions"
import sys, os

import logging
import logging.config
from logging import *
# The python 2.6 version of 'logging' hides these functions, so we need to import explcitly
from logging import getLevelName, root, info, debug, warning, error, critical, getLogger
import ConfigParser

import maya
import pymel.util as util
import maya.utils
from maya.OpenMaya import MGlobal, MEventMessage, MMessage

from pymel.util.decoration import decorator

PYMEL_CONF_ENV_VAR = 'PYMEL_CONF'

#===============================================================================
# DEFAULT FORMAT SETUP
#===============================================================================

def _fixMayaOutput():
    if not hasattr( sys.stdout,"flush"):
        def flush(*args,**kwargs):
            pass
        try:
            sys.stdout.flush = flush
        except AttributeError:
            # second try
            #if hasattr(maya,"Output") and not hasattr(maya.Output,"flush"):
            class MayaOutput(maya.Output):
                def flush(*args,**kwargs):
                    pass
            maya.Output = MayaOutput()
            sys.stdout = maya.Output

_fixMayaOutput()

def getConfigFile():
    if PYMEL_CONF_ENV_VAR in os.environ:
        configFile = os.environ[PYMEL_CONF_ENV_VAR]
        if os.path.isfile(configFile):
            return configFile
    if 'HOME' in os.environ:
        configFile = os.path.join( os.environ['HOME'], "pymel.conf")
        if os.path.isfile(configFile):
            return configFile
    moduleDir = os.path.dirname( os.path.dirname( sys.modules[__name__].__file__ ) )
    configFile = os.path.join(moduleDir,"pymel.conf")
    if os.path.isfile(configFile):
        return configFile
    raise IOError, "Could not find pymel.conf"

def getLogConfigFile():
    configFile = os.path.join(os.path.dirname(__file__),"user_logging.conf")
    if os.path.isfile(configFile):
        return configFile
    return getConfigFile()

assert hasattr(maya.utils, 'shellLogHandler'), "If you manually installed pymel, ensure " \
    "that pymel comes before Maya's site-packages directory on PYTHONPATH / sys.path.  " \
    "See pymel docs for more info."

def pymelLogFileConfig(fname, defaults=None):):
    """
    Reads in a file to set up pymel's loggers.

    Like logging.config.fileConfig, but intended only for pymel's loggers,
    and less heavy handed - fileConfig will, for instance, wipe out all
    handlers from logging._handlers, remove all handlers from logging.root, etc

    For the root logger, the default behavior for pre-existing handlers is to
    keep them - any handlers given in the config file will be 'prepended' to
    the list of root logger handlers. (Maya will set up a 'default' root
    logger, so adding extra root loggers should only be necessary if you have
    specific needs).
    If you wish to to make the setup REMOVE all pre-existing root handlers,
    add a 'remove_existing_handlers' option to the logger_root section,
    and set it's value to True.

    For loggers other than the root logger, the default behavior is the
    opposite - remove all existing handlers; if you wish to keep them,
    set 'remove_existing_handlers' to False.

    The format of the config file read in must meet the same requirements
    as for logging.config.fileConfig - it must have the sections [loggers],
    [handlers], and [fomatters], and it must have an entry for [logger_root]...
    even if not options are set for it.
    """
    cp = ConfigParser.ConfigParser(defaults)
    if hasattr(cp, 'readfp') and hasattr(fname, 'readline'):
        cp.readfp(fname)
    else:
        cp.read(fname)

    formatters = _create_formatters(cp)

    # critical section
    logging._acquireLock()
    try:
        # Handlers add themselves to logging._handlers
        handlers = logging.config._install_handlers(cp, formatters)
        # _install_loggers will remove all existing handlers for the
        # root logger, and any other handlers specified... to override
        # this, save the existing handlers first
        root = logging.root
        rootHandlers = rootHandlers
        oldLogHandlers = {}
        for loggerName, logger in root.manager.loggerDict.iteritems():
            oldLogHandlers[loggerName] = logger.handlers
        if sys.version_info >= (2,6):
            logging.config._install_loggers(configFile, handlers,
                                            disable_existing_loggers=0)
        else:
            logging.config._install_loggers(configFile, handlers)
            # The _install_loggers function disables old-loggers, so we need to
            # re-enable them
            for k,v in logging.root.manager.loggerDict.iteritems():
                if hasattr(v, 'disabled') and v.disabled:
                    v.disabled = 0
        # Now re-add any removed handlers, if needed
        secNames = [cp.get('loggers', 'keys').split(',')]
        secNames = ['logger_' + x.strip() for x in secNames if x != 'root']
        _addOldHandlers(root, rootHandlers, 'logger_root', cp, False )
        for secName in secNames:
            logName = cp.get(secName, 'qualname')
            _addOldHandlers(root.manager.loggerDict[logName],
                            oldLogHandlers[logName], secName, cp, True)
    finally:
        logging._releaseLock()

def _addOldHandlers(logger, oldHandlers, secName, configParser,
                    removeExistingDefault):
    opts = cp.options(secName)
    if 'remove_existing_handlers' not in opts:
        removeExisting = removeExistingDefault
    else:
        removeExisting = eval(opts['remove_existing_handlers'])
    if not removeExisting:
        for handler in oldHandlers:
            if handler not in logger.handlers:
                logger.addhandler(handler)

maya.utils.shellLogHandler()

pymelLogFileConfig(getLogConfigFile())

rootLogger = logging.root

pymelLogger = logging.getLogger("pymel")

# keep as an enumerator so that we can keep the order
logLevels = util.Enum( 'logLevels', dict([(getLevelName(n),n) for n in range(0,CRITICAL+1,10)]) )

def nameToLevel(name):
    return logLevels.getIndex(name)

def levelToName(level):
    return logLevels.getKey(level)

#===============================================================================
# DECORATORS
#===============================================================================

def timed(level=DEBUG):
    import time
    @decorator
    def timedWithLevel(func):
        logger = getLogger(func.__module__)
        def timedFunction(*arg, **kwargs):
            t = time.time()
            res = func(*arg, **kwargs)
            t = time.time() - t # convert to seconds float
            strSecs = time.strftime("%M:%S.", time.localtime(t)) + ("%.3f" % t).split(".")[-1]
            logger.log(level, 'Function %s(...) - finished in %s seconds' % (func.func_name, strSecs))
            return res
        return timedFunction
    return timedWithLevel

#===============================================================================
# INIT TO USER'S PREFERENCE
#===============================================================================

def _setupLevelPreferenceHook():
    """Sets up a callback so that the last used log-level is saved to the user preferences file"""

    LOGLEVEL_OPTVAR = 'pymel.logLevel'

    # retrieve the preference as a string name, for human readability.
    # we need to use MGlobal because cmds.optionVar might not exist yet
    # TODO : resolve load order for standalone.  i don't think that userPrefs is loaded yet at this point in standalone.
    levelName = os.environ.get( 'PYMEL_LOGLEVEL', MGlobal.optionVarStringValue( LOGLEVEL_OPTVAR ) )
    if levelName:
        level =  min( logging.WARNING, nameToLevel(levelName) ) # no more than WARNING level
        pymelLogger.setLevel(level)
        pymelLogger.info("setting logLevel to user preference: %s (%d)" % (levelName, level) )

    func = pymelLogger.setLevel
    def setLevelHook(level, *args, **kwargs):

        levelName = levelToName(level)
        level = nameToLevel(level)
        ret = func(level, *args, **kwargs)
        pymelLogger.info("Log Level Changed to '%s'" % levelName)
        try:
            # save the preference as a string name, for human readability
            # we need to use MGlobal because cmds.optionVar might not exist yet
            MGlobal.setOptionVarValue( LOGLEVEL_OPTVAR, levelName )
        except Exception, e:
            pymelLogger.warning("Log Level could not be saved to the user-prefs ('%s')" % e)
        return ret

    setLevelHook.__doc__ = func.__doc__
    setLevelHook.__name__ = func.__name__
    pymelLogger.setLevel = setLevelHook

#_setupLevelPreferenceHook()

<commit_msg>pymelLogFileConfig fixes<commit_after>"pymel logging functions"
import sys, os

import logging
import logging.config
from logging import *
# The python 2.6 version of 'logging' hides these functions, so we need to import explcitly
from logging import getLevelName, root, info, debug, warning, error, critical, getLogger
import ConfigParser

import maya
import pymel.util as util
import maya.utils
from maya.OpenMaya import MGlobal, MEventMessage, MMessage

from pymel.util.decoration import decorator

PYMEL_CONF_ENV_VAR = 'PYMEL_CONF'

#===============================================================================
# DEFAULT FORMAT SETUP
#===============================================================================

def _fixMayaOutput():
    if not hasattr( sys.stdout,"flush"):
        def flush(*args,**kwargs):
            pass
        try:
            sys.stdout.flush = flush
        except AttributeError:
            # second try
            #if hasattr(maya,"Output") and not hasattr(maya.Output,"flush"):
            class MayaOutput(maya.Output):
                def flush(*args,**kwargs):
                    pass
            maya.Output = MayaOutput()
            sys.stdout = maya.Output

_fixMayaOutput()

def getConfigFile():
    if PYMEL_CONF_ENV_VAR in os.environ:
        configFile = os.environ[PYMEL_CONF_ENV_VAR]
        if os.path.isfile(configFile):
            return configFile
    if 'HOME' in os.environ:
        configFile = os.path.join( os.environ['HOME'], "pymel.conf")
        if os.path.isfile(configFile):
            return configFile
    moduleDir = os.path.dirname( os.path.dirname( sys.modules[__name__].__file__ ) )
    configFile = os.path.join(moduleDir,"pymel.conf")
    if os.path.isfile(configFile):
        return configFile
    raise IOError, "Could not find pymel.conf"

def getLogConfigFile():
    configFile = os.path.join(os.path.dirname(__file__),"user_logging.conf")
    if os.path.isfile(configFile):
        return configFile
    return getConfigFile()

assert hasattr(maya.utils, 'shellLogHandler'), "If you manually installed pymel, ensure " \
    "that pymel comes before Maya's site-packages directory on PYTHONPATH / sys.path.  " \
    "See pymel docs for more info."

#    Like logging.config.fileConfig, but intended only for pymel's loggers,
#    and less heavy handed - fileConfig will, for instance, wipe out all
#    handlers from logging._handlers, remove all handlers from logging.root, etc
def pymelLogFileConfig(fname, defaults=None, disable_existing_loggers=False):
    """
    Reads in a file to set up pymel's loggers.

    In most respects, this function behaves similarly to logging.config.fileConfig -
    consult it's help for details. In particular, the format of the config file
    must meet the same requirements - it must have the sections [loggers],
    [handlers], and [fomatters], and it must have an entry for [logger_root]...
    even if not options are set for it.

    It differs from logging.config.fileConfig in the following ways:

    1) It will not disable any pre-existing loggers which are not specified in
    the config file, unless disable_existing_loggers is set to True.

    2) Like logging.config.fileConfig, the default behavior for pre-existing
    handlers on any loggers whose settings are specified in the config file is
    to remove them; ie, ONLY the handlers explicitly given in the config will
    be on the configured logger.
    However, pymelLogFileConfig provides the ability to keep pre-exisiting
    handlers, by setting the 'remove_existing_handlers' option in the appropriate
    section to True.
    """
    cp = ConfigParser.ConfigParser(defaults)
    if hasattr(cp, 'readfp') and hasattr(fname, 'readline'):
        cp.readfp(fname)
    else:
        cp.read(fname)

    formatters = logging.config._create_formatters(cp)

    # _install_loggers will remove all existing handlers for the
    # root logger, and any other handlers specified... to override
    # this, save the existing handlers first
    root = logging.root
    # make sure you get a COPY of handlers!
    rootHandlers = root.handlers[:]
    oldLogHandlers = {}

    # can't use loggerDict.iteritems, as some of the values are
    # 'PlaceHolder' items... need to use logging.getLogger()
    for loggerName in root.manager.loggerDict:
        logger = logging.getLogger(loggerName)
        # make sure you get a COPY of handlers!
        oldLogHandlers[loggerName] = logger.handlers[:]

    # critical section
    logging._acquireLock()
    try:
        # Handlers add themselves to logging._handlers
        handlers = logging.config._install_handlers(cp, formatters)

        if sys.version_info >= (2,6):
            logging.config._install_loggers(cp, handlers,
                                            disable_existing_loggers=0)
        else:
            logging.config._install_loggers(cp, handlers)
            # The _install_loggers function disables old-loggers, so we need to
            # re-enable them
            for k,v in logging.root.manager.loggerDict.iteritems():
                if hasattr(v, 'disabled') and v.disabled:
                    v.disabled = 0

        # Now re-add any removed handlers, if needed
        secNames = cp.get('loggers', 'keys').split(',')
        secNames = ['logger_' + x.strip() for x in secNames]
        _addOldHandlers(root, rootHandlers, 'logger_root', cp)
        for secName in secNames:
            if secName == 'logger_root':
                logger = root
                oldHandlers = rootHandlers
            else:
                logName = cp.get(secName, 'qualname')
                logger = logging.getLogger(logName)
                oldHandlers = oldLogHandlers.get(logName)
            if oldHandlers:
                _addOldHandlers(logger, oldHandlers, secName, cp)

    finally:
        logging._releaseLock()

def _addOldHandlers(logger, oldHandlers, secName, configParser):
    opts = configParser.options(secName)
    if 'remove_existing_handlers' not in opts:
        removeExisting = True
    else:
        removeExisting = eval(configParser.get(secName, 'remove_existing_handlers'))
    if not removeExisting:
        for handler in oldHandlers:
            if handler not in logger.handlers:
                logger.addHandler(handler)

maya.utils.shellLogHandler()

pymelLogFileConfig(getLogConfigFile())

rootLogger = logging.root

pymelLogger = logging.getLogger("pymel")

# keep as an enumerator so that we can keep the order
logLevels = util.Enum( 'logLevels', dict([(getLevelName(n),n) for n in range(0,CRITICAL+1,10)]) )

def nameToLevel(name):
    return logLevels.getIndex(name)

def levelToName(level):
    return logLevels.getKey(level)

#===============================================================================
# DECORATORS
#===============================================================================

def timed(level=DEBUG):
    import time
    @decorator
    def timedWithLevel(func):
        logger = getLogger(func.__module__)
        def timedFunction(*arg, **kwargs):
            t = time.time()
            res = func(*arg, **kwargs)
            t = time.time() - t # convert to seconds float
            strSecs = time.strftime("%M:%S.", time.localtime(t)) + ("%.3f" % t).split(".")[-1]
            logger.log(level, 'Function %s(...) - finished in %s seconds' % (func.func_name, strSecs))
            return res
        return timedFunction
    return timedWithLevel

#===============================================================================
# INIT TO USER'S PREFERENCE
#===============================================================================

def _setupLevelPreferenceHook():
    """Sets up a callback so that the last used log-level is saved to the user preferences file"""

    LOGLEVEL_OPTVAR = 'pymel.logLevel'

    # retrieve the preference as a string name, for human readability.
    # we need to use MGlobal because cmds.optionVar might not exist yet
    # TODO : resolve load order for standalone.  i don't think that userPrefs is loaded yet at this point in standalone.
    levelName = os.environ.get( 'PYMEL_LOGLEVEL', MGlobal.optionVarStringValue( LOGLEVEL_OPTVAR ) )
    if levelName:
        level =  min( logging.WARNING, nameToLevel(levelName) ) # no more than WARNING level
        pymelLogger.setLevel(level)
        pymelLogger.info("setting logLevel to user preference: %s (%d)" % (levelName, level) )

    func = pymelLogger.setLevel
    def setLevelHook(level, *args, **kwargs):

        levelName = levelToName(level)
        level = nameToLevel(level)
        ret = func(level, *args, **kwargs)
        pymelLogger.info("Log Level Changed to '%s'" % levelName)
        try:
            # save the preference as a string name, for human readability
            # we need to use MGlobal because cmds.optionVar might not exist yet
            MGlobal.setOptionVarValue( LOGLEVEL_OPTVAR, levelName )
        except Exception, e:
            pymelLogger.warning("Log Level could not be saved to the user-prefs ('%s')" % e)
        return ret

    setLevelHook.__doc__ = func.__doc__
    setLevelHook.__name__ = func.__name__
    pymelLogger.setLevel = setLevelHook

#_setupLevelPreferenceHook()

<|endoftext|>


<commit_before>

a = 10<commit_msg>some code clean-ups<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
This module defines structured units and quantities.
"""

# Standard library
import operator

import numpy as np

from .core import Unit, UnitBase

__all__ = ['StructuredUnit']

DTYPE_OBJECT = np.dtype('O')

def _names_from_dtype(dtype):
    """Recursively extract field names from a dtype."""
    names = []
    for name in dtype.names:
        subdtype = dtype.fields[name][0]
        if subdtype.names:
            names.append([name, _names_from_dtype(subdtype)])
        else:
            names.append(name)
    return tuple(names)

def _normalize_names(names):
    """Recursively normalize, inferring upper level names for unadorned tuples.

    Generally, we want the field names to be organized like dtypes, as in
    ``(['pv', ('p', 'v')], 't')``.  But we automatically infer upper
    field names if the list is absent from items like ``(('p', 'v'), 't')``,
    by concatenating the names inside the tuple.
    """
    result = []
    for name in names:
        if isinstance(name, str) and len(name) > 0:
            result.append(name)
        elif (isinstance(name, list)
              and len(name) == 2
              and isinstance(name[0], str) and len(name[0]) > 0
              and isinstance(name[1], tuple) and len(name[1]) > 0):
            result.append([name[0], _normalize_names(name[1])])
        elif isinstance(name, tuple) and len(name) > 0:
            new_tuple = _normalize_names(name)
            result.append([''.join([(i[0] if isinstance(i, list) else i)
                                    for i in new_tuple]), new_tuple])
        else:
            raise ValueError(f'invalid entry {name!r}. Should be a name, '
                             'tuple of names, or 2-element list of the '
                             'form [name, tuple of names].')

    return tuple(result)

class StructuredUnit:
    """Container for units for a structured Quantity.

    Parameters
    ----------
    units : unit-like, tuple of unit-like, or `~astropy.units.StructuredUnit`
        Tuples can be nested.  If a `~astropy.units.StructuredUnit` is passed
        in, it will be returned unchanged unless different names are requested.
    names : tuple of str, tuple or list; `~numpy.dtype`; or `~astropy.units.StructuredUnit`, optional
        Field names for the units, possibly nested. Can be inferred from a
        structured `~numpy.dtype` or another `~astropy.units.StructuredUnit`.
        For nested tuples, by default the name of the upper entry will be the
        concatenation of the names of the lower levels.  One can pass in a
        list with the upper-level name and a tuple of lower-level names to
        avoid this.  For tuples, not all levels have to be given; for any level
        not passed in, default field names of 'f0', 'f1', etc., will be used.

    Notes
    -----
    It is recommended to initialze the class indirectly, using
    `~astropy.units.Unit`.  E.g., ``u.Unit('AU,AU/day')``

    When combined with a structured array to produce a structured
    `~astropy.units.Quantity`, array field names will take precedence.

    """
    def __new__(cls, units, names=None):
        dtype = None
        if names is not None:
            if isinstance(names, StructuredUnit):
                dtype = names._units.dtype
                names = names.field_names
            elif isinstance(names, np.dtype):
                if not names.fields:
                    raise ValueError('dtype should be structured, with fields.')
                dtype = np.dtype([(name, DTYPE_OBJECT) for name in names.names])
                names = _names_from_dtype(names)
            else:
                if not isinstance(names, tuple):
                    names = (names,)
                names = _normalize_names(names)

        if not isinstance(units, tuple):
            units = Unit(units)
            if isinstance(units, StructuredUnit):
                # Avoid constructing a new StructuredUnit if no field names
                # are given, or if all field names are the same already anyway.
                if names is None or units.field_names == names:
                    return units

                # Otherwise, turn (the upper level) into a tuple, for renaming.
                units = units.values()
            else:
                # Single regular unit: make a tuple for iteration below.
                units = (units,)

        if names is None:
            names = tuple(f'f{i}' for i in range(len(units)))

        elif len(units) != len(names):
            raise ValueError("lengths of units and field names must match.")

        converted = []
        for unit, name in zip(units, names):
            if isinstance(name, list):
                # For list, the first item is the name of our level,
                # and the second another tuple of names, i.e., we recurse.
                unit = cls(unit, name[1])
                name = name[0]
            else:
                # We are at the lowest level.  Check unit.
                unit = Unit(unit)
                if dtype is not None and isinstance(unit, StructuredUnit):
                    raise ValueError("units do not match in depth with field "
                                     "names from dtype or structured unit.")

            converted.append(unit)

        self = super().__new__(cls)
        if dtype is None:
            dtype = np.dtype([((name[0] if isinstance(name, list) else name),
                               DTYPE_OBJECT) for name in names])
        # Decay array to void so we can access by field name and number.
        self._units = np.array(tuple(converted), dtype)[()]
        return self

    @property
    def field_names(self):
        """Possibly nested tuple of the field names of the parts."""
        return tuple(([name, unit.field_names]
                      if isinstance(unit, StructuredUnit) else name)
                     for name, unit in self.items())

    # Allow StructuredUnit to be treated as an (ordered) mapping.
    def __len__(self):
        return len(self._units.dtype.names)

    def __getitem__(self, item):
        # Since we are based on np.void, indexing by field number works too.
        return self._units[item]

    def values(self):
        return self._units.item()

    def keys(self):
        return self._units.dtype.names

    def items(self):
        return zip(self._units.dtype.names, self._units.item())

    def __iter__(self):
        yield from self._units.dtype.names

    # Helpers for methods below.
    def _recursively_apply(self, func, cls=None):
        """Apply func recursively.

        Parameters
        ----------
        func : callable
            Function to apply to all parts of the structured unit,
            recursing as needed.
        cls : type, optional
            If given, should be a subclass of `~numpy.void`. By default,
            will return a new `~astropy.units.StructuredUnit` instance.
        """
        results = np.array(tuple([func(part) for part in self.values()]),
                           self._units.dtype)[()]
        if cls is not None:
            return results.view((cls, results.dtype))

        # Short-cut; no need to interpret field names, etc.
        result = super().__new__(self.__class__)
        result._units = results
        return result

    def _recursively_get_dtype(self, value, enter_lists=True):
        """Get structured dtype according to value, using our field names.

        This is useful since ``np.array(value)`` would treat tuples as lower
        levels of the array, rather than as elements of a structured array.
        The routine does presume that the type of the first tuple is
        representative of the rest.  Used in ``_get_converter``.

        """
        if enter_lists:
            while isinstance(value, list):
                value = value[0]
        if not isinstance(value, tuple) or len(self) != len(value):
            raise ValueError(f"cannot interpret value {value} for unit {self}.")
        descr = []
        for (name, unit), part in zip(self.items(), value):
            if isinstance(unit, StructuredUnit):
                descr.append(
                    (name, unit._recursively_get_dtype(part, enter_lists=False)))
            else:
                # Got a part associated with a regular unit. Gets its dtype.
                # Like for Quantity, we cast integers to float.
                part = np.array(part)
                part_dtype = part.dtype
                if part_dtype.kind in 'iu':
                    part_dtype = np.dtype(float)
                descr.append((name, part_dtype, part.shape))
        return np.dtype(descr)

    @property
    def si(self):
        """The `StructuredUnit` instance in SI units."""
        return self._recursively_apply(operator.attrgetter('si'))

    @property
    def cgs(self):
        """The `StructuredUnit` instance in cgs units."""
        return self._recursively_apply(operator.attrgetter('cgs'))

    # Needed to pass through Unit initializer, so might as well use it.
    def _get_physical_type_id(self):
        return self._recursively_apply(
            operator.methodcaller('_get_physical_type_id'), cls=Structure)

    @property
    def physical_type(self):
        """Physical types of all the fields."""
        return self._recursively_apply(
            operator.attrgetter('physical_type'), cls=Structure)

    def decompose(self, bases=set()):
        """The `StructuredUnit` composed of only irreducible units.

        Parameters
        ----------
        bases : sequence of `~astropy.units.UnitBase`, optional
            The bases to decompose into.  When not provided,
            decomposes down to any irreducible units.  When provided,
            the decomposed result will only contain the given units.
            This will raises a `UnitsError` if it's not possible
            to do so.

        Returns
        -------
        `~astropy.units.StructuredUnit`
            With the unit for each field containing only irreducible units.
        """
        return self._recursively_apply(
            operator.methodcaller('decompose', bases=bases))

    def is_equivalent(self, other, equivalencies=[]):
        """`True` if all fields are equivalent to the other's fields.

        Parameters
        ----------
        other : `~astropy.units.StructuredUnit`
            The structured unit to compare with, or what can initialize one.
        equivalencies : list of tuple, optional
            A list of equivalence pairs to try if the units are not
            directly convertible.  See :ref:`unit_equivalencies`.
            The list will be applied to all fields.

        Returns
        -------
        bool
        """
        try:
            other = StructuredUnit(other)
        except Exception:
            return False

        if len(self) != len(other):
            return False

        for self_part, other_part in zip(self.values(), other.values()):
            if not self_part.is_equivalent(other_part,
                                           equivalencies=equivalencies):
                return False

        return True

    def _get_converter(self, other, equivalencies=[]):
        if not isinstance(other, type(self)):
            other = self.__class__(other, names=self)

        converters = [self_part._get_converter(other_part,
                                               equivalencies=equivalencies)
                      for (self_part, other_part) in zip(self.values(),
                                                         other.values())]

        def converter(value):
            if not hasattr(value, 'dtype'):
                value = np.array(value, self._recursively_get_dtype(value))
            result = np.empty_like(value)
            for name, converter_ in zip(result.dtype.names, converters):
                result[name] = converter_(value[name])
            return result

        return converter

    def to(self, other, value, equivalencies=[]):
        """Return values converted to the specified unit.

        Parameters
        ----------
        other : `~astropy.units.StructuredUnit`
            The unit to convert to.  If necessary, will be converted to
            a `~astropy.units.StructuredUnit` using the dtype of ``value``.
        value : array-like
            Value(s) in the current unit to be converted to the
            specified unit.  If a sequence, the first element must have
            entries of the correct type to represent all elements (i.e.,
            not have, e.g., a ``float`` where other elements have ``complex``).
        equivalencies : list of tuple, optional
            A list of equivalence pairs to try if the units are not
            directly convertible.  See :ref:`unit_equivalencies`.
            This list is in addition to possible global defaults set by, e.g.,
            `set_enabled_equivalencies`.
            Use `None` to turn off all equivalencies.

        Returns
        -------
        values : scalar or array
            Converted value(s).

        Raises
        ------
        UnitsError
            If units are inconsistent
        """
        return self._get_converter(other, equivalencies=equivalencies)(value)

    def to_string(self, format='generic'):
        """Output the unit in the given format as a string.

        Units are separated by commas.

        Parameters
        ----------
        format : `astropy.units.format.Base` instance or str
            The name of a format or a formatter object.  If not
            provided, defaults to the generic format.

        Notes
        -----
        Structured units can be written to all formats, but can be
        re-read only with 'generic'.

        """
        parts = [part.to_string(format) for part in self.values()]
        out_fmt = '({})' if len(self) > 1 else '({},)'
        if format == 'latex':
            # Strip $ from parts and add them on the outside.
            parts = [part[1:-1] for part in parts]
            out_fmt = '$' + out_fmt + '$'
        return out_fmt.format(', '.join(parts))

    def _repr_latex_(self):
        return self.to_string('latex')

    __array_ufunc__ = None

    def __mul__(self, other):
        if isinstance(other, str):
            try:
                other = Unit(other, parse_strict='silent')
            except Exception:
                return NotImplemented
        if isinstance(other, UnitBase):
            new_units = tuple(part * other for part in self.values())
            return self.__class__(new_units, names=self)
        if isinstance(other, StructuredUnit):
            return NotImplemented

        # Anything not like a unit, try initialising as a structured quantity.
        try:
            from .quantity import Quantity
            return Quantity(other, unit=self)
        except Exception:
            return NotImplemented

    def __rmul__(self, other):
        return self.__mul__(other)

    def __truediv__(self, other):
        if isinstance(other, str):
            try:
                other = Unit(other, parse_strict='silent')
            except Exception:
                return NotImplemented

        if isinstance(other, UnitBase):
            new_units = tuple(part / other for part in self.values())
            return self.__class__(new_units, names=self)
        return NotImplemented

    def __rlshift__(self, m):
        try:
            from .quantity import Quantity
            return Quantity(m, self, copy=False, subok=True)
        except Exception:
            return NotImplemented

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

    def __repr__(self):
        return f'Unit("{self.to_string()}")'

    def __eq__(self, other):
        try:
            other = StructuredUnit(other)
        except Exception:
            return NotImplemented

        return self.values() == other.values()

    def __ne__(self, other):
        if not isinstance(other, type(self)):
            try:
                other = StructuredUnit(other)
            except Exception:
                return NotImplemented

        return self.values() != other.values()

class Structure(np.void):
    """Single element structure for physical type IDs, etc.

    Behaves like a `~numpy.void` and thus mostly like a tuple which can also
    be indexed with field names, but overrides ``__eq__`` and ``__ne__`` to
    compare only the contents, not the field names are ignored.  Furthermore,
    this way no `FutureWarning` about comparisons is given.

    """
    # Note that it is important that it is important for physical type IDs to
    # not be stored in a tuple, since then the physical types would be treated
    # as alternatives in :meth:`~astropy.units.UnitBase.is_equivalent`.
    # (Of course, in that case, they could also not be indexed by name.)

    def __eq__(self, other):
        if isinstance(other, np.void):
            other = other.item()

        return self.item() == other

    def __ne__(self, other):
        if isinstance(other, np.void):
            other = other.item()

        return self.item() != other
<commit_msg>Make .to(other) return the result for a value with all ones.<commit_after># -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
This module defines structured units and quantities.
"""

# Standard library
import operator

import numpy as np

from .core import Unit, UnitBase, UNITY

__all__ = ['StructuredUnit']

DTYPE_OBJECT = np.dtype('O')

def _names_from_dtype(dtype):
    """Recursively extract field names from a dtype."""
    names = []
    for name in dtype.names:
        subdtype = dtype.fields[name][0]
        if subdtype.names:
            names.append([name, _names_from_dtype(subdtype)])
        else:
            names.append(name)
    return tuple(names)

def _normalize_names(names):
    """Recursively normalize, inferring upper level names for unadorned tuples.

    Generally, we want the field names to be organized like dtypes, as in
    ``(['pv', ('p', 'v')], 't')``.  But we automatically infer upper
    field names if the list is absent from items like ``(('p', 'v'), 't')``,
    by concatenating the names inside the tuple.
    """
    result = []
    for name in names:
        if isinstance(name, str) and len(name) > 0:
            result.append(name)
        elif (isinstance(name, list)
              and len(name) == 2
              and isinstance(name[0], str) and len(name[0]) > 0
              and isinstance(name[1], tuple) and len(name[1]) > 0):
            result.append([name[0], _normalize_names(name[1])])
        elif isinstance(name, tuple) and len(name) > 0:
            new_tuple = _normalize_names(name)
            result.append([''.join([(i[0] if isinstance(i, list) else i)
                                    for i in new_tuple]), new_tuple])
        else:
            raise ValueError(f'invalid entry {name!r}. Should be a name, '
                             'tuple of names, or 2-element list of the '
                             'form [name, tuple of names].')

    return tuple(result)

class StructuredUnit:
    """Container for units for a structured Quantity.

    Parameters
    ----------
    units : unit-like, tuple of unit-like, or `~astropy.units.StructuredUnit`
        Tuples can be nested.  If a `~astropy.units.StructuredUnit` is passed
        in, it will be returned unchanged unless different names are requested.
    names : tuple of str, tuple or list; `~numpy.dtype`; or `~astropy.units.StructuredUnit`, optional
        Field names for the units, possibly nested. Can be inferred from a
        structured `~numpy.dtype` or another `~astropy.units.StructuredUnit`.
        For nested tuples, by default the name of the upper entry will be the
        concatenation of the names of the lower levels.  One can pass in a
        list with the upper-level name and a tuple of lower-level names to
        avoid this.  For tuples, not all levels have to be given; for any level
        not passed in, default field names of 'f0', 'f1', etc., will be used.

    Notes
    -----
    It is recommended to initialze the class indirectly, using
    `~astropy.units.Unit`.  E.g., ``u.Unit('AU,AU/day')``

    When combined with a structured array to produce a structured
    `~astropy.units.Quantity`, array field names will take precedence.

    """
    def __new__(cls, units, names=None):
        dtype = None
        if names is not None:
            if isinstance(names, StructuredUnit):
                dtype = names._units.dtype
                names = names.field_names
            elif isinstance(names, np.dtype):
                if not names.fields:
                    raise ValueError('dtype should be structured, with fields.')
                dtype = np.dtype([(name, DTYPE_OBJECT) for name in names.names])
                names = _names_from_dtype(names)
            else:
                if not isinstance(names, tuple):
                    names = (names,)
                names = _normalize_names(names)

        if not isinstance(units, tuple):
            units = Unit(units)
            if isinstance(units, StructuredUnit):
                # Avoid constructing a new StructuredUnit if no field names
                # are given, or if all field names are the same already anyway.
                if names is None or units.field_names == names:
                    return units

                # Otherwise, turn (the upper level) into a tuple, for renaming.
                units = units.values()
            else:
                # Single regular unit: make a tuple for iteration below.
                units = (units,)

        if names is None:
            names = tuple(f'f{i}' for i in range(len(units)))

        elif len(units) != len(names):
            raise ValueError("lengths of units and field names must match.")

        converted = []
        for unit, name in zip(units, names):
            if isinstance(name, list):
                # For list, the first item is the name of our level,
                # and the second another tuple of names, i.e., we recurse.
                unit = cls(unit, name[1])
                name = name[0]
            else:
                # We are at the lowest level.  Check unit.
                unit = Unit(unit)
                if dtype is not None and isinstance(unit, StructuredUnit):
                    raise ValueError("units do not match in depth with field "
                                     "names from dtype or structured unit.")

            converted.append(unit)

        self = super().__new__(cls)
        if dtype is None:
            dtype = np.dtype([((name[0] if isinstance(name, list) else name),
                               DTYPE_OBJECT) for name in names])
        # Decay array to void so we can access by field name and number.
        self._units = np.array(tuple(converted), dtype)[()]
        return self

    @property
    def field_names(self):
        """Possibly nested tuple of the field names of the parts."""
        return tuple(([name, unit.field_names]
                      if isinstance(unit, StructuredUnit) else name)
                     for name, unit in self.items())

    # Allow StructuredUnit to be treated as an (ordered) mapping.
    def __len__(self):
        return len(self._units.dtype.names)

    def __getitem__(self, item):
        # Since we are based on np.void, indexing by field number works too.
        return self._units[item]

    def values(self):
        return self._units.item()

    def keys(self):
        return self._units.dtype.names

    def items(self):
        return tuple(zip(self._units.dtype.names, self._units.item()))

    def __iter__(self):
        yield from self._units.dtype.names

    # Helpers for methods below.
    def _recursively_apply(self, func, cls=None):
        """Apply func recursively.

        Parameters
        ----------
        func : callable
            Function to apply to all parts of the structured unit,
            recursing as needed.
        cls : type, optional
            If given, should be a subclass of `~numpy.void`. By default,
            will return a new `~astropy.units.StructuredUnit` instance.
        """
        results = np.array(tuple([func(part) for part in self.values()]),
                           self._units.dtype)[()]
        if cls is not None:
            return results.view((cls, results.dtype))

        # Short-cut; no need to interpret field names, etc.
        result = super().__new__(self.__class__)
        result._units = results
        return result

    def _recursively_get_dtype(self, value, enter_lists=True):
        """Get structured dtype according to value, using our field names.

        This is useful since ``np.array(value)`` would treat tuples as lower
        levels of the array, rather than as elements of a structured array.
        The routine does presume that the type of the first tuple is
        representative of the rest.  Used in ``_get_converter``.

        For the special value of ``UNITY``, all fields are assumed to be 1.0,
        and hence this will return an all-float dtype.

        """
        if enter_lists:
            while isinstance(value, list):
                value = value[0]
        if value is UNITY:
            value = (UNITY,) * len(self)
        elif not isinstance(value, tuple) or len(self) != len(value):
            raise ValueError(f"cannot interpret value {value} for unit {self}.")
        descr = []
        for (name, unit), part in zip(self.items(), value):
            if isinstance(unit, StructuredUnit):
                descr.append(
                    (name, unit._recursively_get_dtype(part, enter_lists=False)))
            else:
                # Got a part associated with a regular unit. Gets its dtype.
                # Like for Quantity, we cast integers to float.
                part = np.array(part)
                part_dtype = part.dtype
                if part_dtype.kind in 'iu':
                    part_dtype = np.dtype(float)
                descr.append((name, part_dtype, part.shape))
        return np.dtype(descr)

    @property
    def si(self):
        """The `StructuredUnit` instance in SI units."""
        return self._recursively_apply(operator.attrgetter('si'))

    @property
    def cgs(self):
        """The `StructuredUnit` instance in cgs units."""
        return self._recursively_apply(operator.attrgetter('cgs'))

    # Needed to pass through Unit initializer, so might as well use it.
    def _get_physical_type_id(self):
        return self._recursively_apply(
            operator.methodcaller('_get_physical_type_id'), cls=Structure)

    @property
    def physical_type(self):
        """Physical types of all the fields."""
        return self._recursively_apply(
            operator.attrgetter('physical_type'), cls=Structure)

    def decompose(self, bases=set()):
        """The `StructuredUnit` composed of only irreducible units.

        Parameters
        ----------
        bases : sequence of `~astropy.units.UnitBase`, optional
            The bases to decompose into.  When not provided,
            decomposes down to any irreducible units.  When provided,
            the decomposed result will only contain the given units.
            This will raises a `UnitsError` if it's not possible
            to do so.

        Returns
        -------
        `~astropy.units.StructuredUnit`
            With the unit for each field containing only irreducible units.
        """
        return self._recursively_apply(
            operator.methodcaller('decompose', bases=bases))

    def is_equivalent(self, other, equivalencies=[]):
        """`True` if all fields are equivalent to the other's fields.

        Parameters
        ----------
        other : `~astropy.units.StructuredUnit`
            The structured unit to compare with, or what can initialize one.
        equivalencies : list of tuple, optional
            A list of equivalence pairs to try if the units are not
            directly convertible.  See :ref:`unit_equivalencies`.
            The list will be applied to all fields.

        Returns
        -------
        bool
        """
        try:
            other = StructuredUnit(other)
        except Exception:
            return False

        if len(self) != len(other):
            return False

        for self_part, other_part in zip(self.values(), other.values()):
            if not self_part.is_equivalent(other_part,
                                           equivalencies=equivalencies):
                return False

        return True

    def _get_converter(self, other, equivalencies=[]):
        if not isinstance(other, type(self)):
            other = self.__class__(other, names=self)

        converters = [self_part._get_converter(other_part,
                                               equivalencies=equivalencies)
                      for (self_part, other_part) in zip(self.values(),
                                                         other.values())]

        def converter(value):
            if not hasattr(value, 'dtype'):
                value = np.array(value, self._recursively_get_dtype(value))
            result = np.empty_like(value)
            for name, converter_ in zip(result.dtype.names, converters):
                result[name] = converter_(value[name])
            # Index with empty tuple to decay array scalars to numpy void.
            return result if result.shape else result[()]

        return converter

    def to(self, other, value=np._NoValue, equivalencies=[]):
        """Return values converted to the specified unit.

        Parameters
        ----------
        other : `~astropy.units.StructuredUnit`
            The unit to convert to.  If necessary, will be converted to
            a `~astropy.units.StructuredUnit` using the dtype of ``value``.
        value : array-like, optional
            Value(s) in the current unit to be converted to the
            specified unit.  If a sequence, the first element must have
            entries of the correct type to represent all elements (i.e.,
            not have, e.g., a ``float`` where other elements have ``complex``).
            If not given, assumed to have 1. in all fields.
        equivalencies : list of tuple, optional
            A list of equivalence pairs to try if the units are not
            directly convertible.  See :ref:`unit_equivalencies`.
            This list is in addition to possible global defaults set by, e.g.,
            `set_enabled_equivalencies`.
            Use `None` to turn off all equivalencies.

        Returns
        -------
        values : scalar or array
            Converted value(s).

        Raises
        ------
        UnitsError
            If units are inconsistent
        """
        if value is np._NoValue:
            # We do not have UNITY as a default, since then the docstring
            # would list 1.0 as default, yet one could not pass that in.
            value = UNITY
        return self._get_converter(other, equivalencies=equivalencies)(value)

    def to_string(self, format='generic'):
        """Output the unit in the given format as a string.

        Units are separated by commas.

        Parameters
        ----------
        format : `astropy.units.format.Base` instance or str
            The name of a format or a formatter object.  If not
            provided, defaults to the generic format.

        Notes
        -----
        Structured units can be written to all formats, but can be
        re-read only with 'generic'.

        """
        parts = [part.to_string(format) for part in self.values()]
        out_fmt = '({})' if len(self) > 1 else '({},)'
        if format == 'latex':
            # Strip $ from parts and add them on the outside.
            parts = [part[1:-1] for part in parts]
            out_fmt = '$' + out_fmt + '$'
        return out_fmt.format(', '.join(parts))

    def _repr_latex_(self):
        return self.to_string('latex')

    __array_ufunc__ = None

    def __mul__(self, other):
        if isinstance(other, str):
            try:
                other = Unit(other, parse_strict='silent')
            except Exception:
                return NotImplemented
        if isinstance(other, UnitBase):
            new_units = tuple(part * other for part in self.values())
            return self.__class__(new_units, names=self)
        if isinstance(other, StructuredUnit):
            return NotImplemented

        # Anything not like a unit, try initialising as a structured quantity.
        try:
            from .quantity import Quantity
            return Quantity(other, unit=self)
        except Exception:
            return NotImplemented

    def __rmul__(self, other):
        return self.__mul__(other)

    def __truediv__(self, other):
        if isinstance(other, str):
            try:
                other = Unit(other, parse_strict='silent')
            except Exception:
                return NotImplemented

        if isinstance(other, UnitBase):
            new_units = tuple(part / other for part in self.values())
            return self.__class__(new_units, names=self)
        return NotImplemented

    def __rlshift__(self, m):
        try:
            from .quantity import Quantity
            return Quantity(m, self, copy=False, subok=True)
        except Exception:
            return NotImplemented

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

    def __repr__(self):
        return f'Unit("{self.to_string()}")'

    def __eq__(self, other):
        try:
            other = StructuredUnit(other)
        except Exception:
            return NotImplemented

        return self.values() == other.values()

    def __ne__(self, other):
        if not isinstance(other, type(self)):
            try:
                other = StructuredUnit(other)
            except Exception:
                return NotImplemented

        return self.values() != other.values()

class Structure(np.void):
    """Single element structure for physical type IDs, etc.

    Behaves like a `~numpy.void` and thus mostly like a tuple which can also
    be indexed with field names, but overrides ``__eq__`` and ``__ne__`` to
    compare only the contents, not the field names are ignored.  Furthermore,
    this way no `FutureWarning` about comparisons is given.

    """
    # Note that it is important that it is important for physical type IDs to
    # not be stored in a tuple, since then the physical types would be treated
    # as alternatives in :meth:`~astropy.units.UnitBase.is_equivalent`.
    # (Of course, in that case, they could also not be indexed by name.)

    def __eq__(self, other):
        if isinstance(other, np.void):
            other = other.item()

        return self.item() == other

    def __ne__(self, other):
        if isinstance(other, np.void):
            other = other.item()

        return self.item() != other
<|endoftext|>