Added UTC conversion function to utils.py

pywwt/core.py

@@ -13,10 +13,10 @@
 from .solar_system import SolarSystem
 from .layers import LayerManager
 from .instruments import Instruments
+from .utils import ensure_utc
 
 import json
 import os
-import pytz
 import shutil
 import tempfile
 
@@ -229,21 +229,8 @@ def set_current_time(self, dt=None):
             uses the current time
         """
         # Ensure the object received is a datetime or Time; convert it to UTC
-        if dt is None:
-            utc_dt = datetime.utcnow().astimezone(pytz.UTC).isoformat()
-        elif isinstance(dt, datetime):
-            if dt.tzinfo is None:
-                utc_dt = pytz.utc.localize(dt).isoformat()
-            elif dt.tzinfo == pytz.UTC:
-                utc_dt = dt.isoformat()
-            else:  # has a non-UTC time zone
-                utc_dt = dt.astimezone(pytz.UTC).isoformat()
-        elif isinstance(dt, Time):
-            utc_dt = dt.to_datetime(pytz.UTC).isoformat()
-        else:
-            raise ValueError('Time must be a datetime or astropy.Time object')
-
-        self._send_msg(event='set_datetime', isot=utc_dt)
+        utc_tm = ensure_utc(dt, str_allowed=False)
+        self._send_msg(event='set_datetime', isot=utc_tm)
 
     def center_on_coordinates(self, coord, fov=60 * u.deg, instant=True):
         """

pywwt/layers.py

@@ -4,7 +4,6 @@
 from os import path
 import shutil
 
-import pytz
 import re
 
 if sys.version_info[0] == 2:  # noqa
@@ -26,7 +25,7 @@
 
 from traitlets import HasTraits, validate, observe
 from .traits import Color, Bool, Float, Unicode, AstropyQuantity, Any, to_hex
-from .utils import sanitize_image, validate_traits
+from .utils import sanitize_image, validate_traits, ensure_utc
 
 __all__ = ['LayerManager', 'TableLayer', 'ImageLayer']
 
@@ -707,39 +706,12 @@ def _on_time_att_change(self, *value):
                                   setting='startDateColumn', value=-1)
             return
 
-        # Convert time column to UTC so WWT displays points at expected times
         wwt_times = Column(self.table[self.time_att].copy()).tolist()
         # must specify Column so we can use tolist() on astropy Time columns
 
+        # Convert time column to UTC so WWT displays points at expected times
         for i, tm in enumerate(wwt_times):
-            if isinstance(tm, datetime):
-                if tm.tzinfo is None:
-                    wwt_times[i] = pytz.utc.localize(tm).isoformat()
-                elif tm.tzinfo == pytz.UTC:
-                    wwt_times[i] = tm.isoformat()
-                else: # has another time zone besides UTC
-                    wwt_times[i] = tm.astimezone(pytz.UTC).isoformat()
-
-            elif isinstance(tm, Time):
-                wwt_times[i] = tm.to_datetime(pytz.UTC).isoformat()
-
-            else: # is an ISOT string
-                utc_tm = Time(tm, format='isot').to_datetime(pytz.UTC)
-                wwt_times[i] = utc_tm.isoformat()
-
-        '''
-        col = self.table[self.time_att]
-
-        if isinstance(col, datetime):
-            wwt_times = Column([t.astimezone(timezone=pytz.UTC).isoformat()
-                                for t in col])
-        elif isinstance(col, Time):
-            wwt_times = Column([t.to_datetime(timezone=pytz.UTC).isot
-                                for t in col])
-        else:
-            col_list = col.tolist()
-            wwt_times = Column(col_list)
-        '''
+            wwt_times[i] = ensure_utc(tm, str_allowed=True)
 
         # Update the table passed to WWT with the new, modified time column
         self.table[TIME_COLUMN_NAME] = Column(wwt_times)
@@ -820,8 +792,6 @@ def _on_trait_change(self, changed):
                 value = VALID_LON_UNITS[self._check_lon_unit({'value': value})]
             elif changed['name'] == 'xyz_unit':
                 value = VALID_ALT_UNITS[self._check_xyz_unit({'value': value})]
-            elif isinstance(value, u.Quantity):
-                value = value.value
             elif changed['name'] == 'time_decay':
                 value = value.to(u.day).value
             self.parent._send_msg(event='table_layer_set',

pywwt/utils.py

@@ -1,6 +1,9 @@
 import numpy as np
+import pytz
 from astropy.io import fits
 from astropy.coordinates import ICRS
+from astropy.time import Time
+from datetime import datetime
 from reproject import reproject_interp
 from reproject.mosaicking import find_optimal_celestial_wcs
 
@@ -31,11 +34,41 @@ def sanitize_image(image, output_file, overwrite=False):
 def validate_traits(cls, traits):
     '''
     Helper function to ensure user-provided trait names match those of the
-    class they're being used to instantiate
+    class they're being used to instantiate.
     '''
     mismatch = [key for key in traits if key not in cls.trait_names()]
     if mismatch:
         raise KeyError('Key{0} {1} do{2}n\'t match any layer trait name'
                        .format('s' if len(mismatch) > 1 else '',
                                mismatch,
                                '' if len(mismatch) > 1 else 'es'))
+
+def ensure_utc(tm, str_allowed):
+    '''
+    Helper function to convert a time object (Time, datetime, or UTC string
+    if str_allowed == True) into UTC before passing it to WWT.
+    str_allowed is True for wwt.set_current_time (core.py) and False for TableLayer's 'time_att' implementation (layers.py).
+    '''
+
+    if tm is None:
+        utc_tm = datetime.utcnow().astimezone(pytz.UTC).isoformat()
+
+    elif isinstance(tm, datetime):
+        if tm.tzinfo is None:
+            utc_tm = pytz.utc.localize(tm).isoformat()
+        elif tm.tzinfo == pytz.UTC:
+            utc_tm = tm.isoformat()
+        else:  # has a non-UTC time zone
+            utc_tm = tm.astimezone(pytz.UTC).isoformat()
+
+    elif isinstance(tm, Time):
+        utc_tm = tm.to_datetime(pytz.UTC).isoformat()
+
+    else:
+        if str_allowed: # is an ISOT string
+            dt = Time(tm, format='isot').to_datetime(pytz.UTC)
+            utc_tm = dt.isoformat()
+        else:
+            raise ValueError('Time must be a datetime or astropy.Time object')
+
+    return utc_tm