Use download() in openmc-get-photon-data and respond to other comments

openmc/_utils.py

@@ -31,31 +31,32 @@ def download(url, checksum=None, as_browser=False, **kwargs):
         page = Request(url, headers={'User-Agent': 'Mozilla/5.0'})
     else:
         page = url
-    req = urlopen(page, **kwargs)
-    # Get file size from header
-    file_size = req.length
 
-    # Check if file already downloaded
-    basename = Path(urlparse(url).path).name
-    if os.path.exists(basename):
-        if os.path.getsize(basename) == file_size:
-            print('Skipping {}, already downloaded'.format(basename))
-            return basename
+    with urlopen(page, **kwargs) as response:
+        # Get file size from header
+        file_size = response.length
 
-    # Copy file to disk in chunks
-    print('Downloading {}... '.format(basename), end='')
-    downloaded = 0
-    with open(basename, 'wb') as fh:
-        while True:
-            chunk = req.read(_BLOCK_SIZE)
-            if not chunk:
-                break
-            fh.write(chunk)
-            downloaded += len(chunk)
-            status = '{:10}  [{:3.2f}%]'.format(
-                downloaded, downloaded * 100. / file_size)
-            print(status + '\b'*len(status), end='')
-        print('')
+        # Check if file already downloaded
+        basename = Path(urlparse(url).path).name
+        if os.path.exists(basename):
+            if os.path.getsize(basename) == file_size:
+                print('Skipping {}, already downloaded'.format(basename))
+                return basename
+
+        # Copy file to disk in chunks
+        print('Downloading {}... '.format(basename), end='')
+        downloaded = 0
+        with open(basename, 'wb') as fh:
+            while True:
+                chunk = response.read(_BLOCK_SIZE)
+                if not chunk:
+                    break
+                fh.write(chunk)
+                downloaded += len(chunk)
+                status = '{:10}  [{:3.2f}%]'.format(
+                    downloaded, downloaded * 100. / file_size)
+                print(status + '\b'*len(status), end='', flush=True)
+            print('')
 
     if checksum is not None:
         downloadsum = hashlib.md5(open(basename, 'rb').read()).hexdigest()

scripts/openmc-get-photon-data

@@ -6,21 +6,20 @@ relaxation data and convert it to an HDF5 library for use with OpenMC.
 This data is used for photon transport in OpenMC.
 """
 
-import os
-import sys
-import shutil
-import zipfile
 import argparse
-from io import BytesIO
-from urllib.request import urlopen
+import os
+from pathlib import Path
+import zipfile
 
 import openmc.data
+from openmc._utils import download
 
 
 class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
                       argparse.RawDescriptionHelpFormatter):
     pass
 
+
 parser = argparse.ArgumentParser(
     description=__doc__,
     formatter_class=CustomFormatter
@@ -36,47 +35,17 @@ block_size = 16384
 # ==============================================================================
 # DOWNLOAD FILES FROM NNDC SITE
 
-if not os.path.exists('photon_hdf5'):
-    os.mkdir('photon_hdf5')
+output = Path('photon_hdf5')
+output.mkdir(exist_ok=True)
 
 for f in files:
-    # Establish connection to URL
-    url = base_url + f
-    req = urlopen(url)
-
-    # Get file size from header
-    file_size = req.length
-    downloaded = 0
-
-    # Check if file already downloaded
-    if os.path.exists(f):
-        if os.path.getsize(f) == file_size:
-            print('Skipping ' + f)
-            continue
-        else:
-            overwrite = input('Overwrite {}? ([y]/n) '.format(f))
-            if overwrite.lower().startswith('n'):
-                continue
-
-    # Copy file to disk
-    print('Downloading {}... '.format(f), end='')
-    with open(f, 'wb') as fh:
-        while True:
-            chunk = req.read(block_size)
-            if not chunk:
-                break
-            fh.write(chunk)
-            downloaded += len(chunk)
-            status = '{0:10}  [{1:3.2f}%]'.format(
-                downloaded, downloaded * 100. / file_size)
-            print(status + chr(8)*len(status), end='')
-        print('')
+    download(base_url + f)
 
 # ==============================================================================
 # EXTRACT FILES
 
 for f in files:
-    print('Extracting {0}...'.format(f))
+    print('Extracting {}...'.format(f))
     zipfile.ZipFile(f).extractall()
 
 # ==============================================================================
@@ -87,9 +56,10 @@ if args.cross_sections is not None:
     lib_path = args.cross_sections
     library = openmc.data.DataLibrary.from_xml(lib_path)
 else:
-    lib_path = os.path.join('photon_hdf5', 'cross_sections.xml')
+    lib_path = output / 'cross_sections.xml'
     library = openmc.data.DataLibrary()
 
+# Iterate over each natural element from Z=1 to Z=100
 for z in range(1, 101):
     element = openmc.data.ATOMIC_SYMBOL[z]
     print('Generating HDF5 file for Z={} ({})...'.format(z, element))
@@ -97,11 +67,11 @@ for z in range(1, 101):
     # Generate instance of IncidentPhoton
     photo_file = os.path.join('photoat', 'photoat-{:03}_{}_000.endf'.format(z, element))
     atom_file = os.path.join('atomic_relax', 'atom-{:03}_{}_000.endf'.format(z, element))
-    f = openmc.data.IncidentPhoton.from_endf(photo_file, atom_file)
+    data = openmc.data.IncidentPhoton.from_endf(photo_file, atom_file)
 
     # Write HDF5 file and register it
-    hdf5_file = os.path.join('photon_hdf5', element + '.h5')
-    f.export_to_hdf5(hdf5_file, 'w')
+    hdf5_file = output / (element + '.h5')
+    data.export_to_hdf5(hdf5_file, 'w')
     library.register_file(hdf5_file)
 
 library.export_to_xml(lib_path)