o Adding python2 compatibility.

4 files changed, 36 insertions, 13 deletions

diff --git a/.gitignore b/.gitignore
index b7e58a7..65e9278 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,8 @@
 *.pyc
 env
 
+*.png
+
 # LTSpice
 demo/*/*.net
 demo/*/*.raw
diff --git a/README.md b/README.md
index 1bcc212..74f98a9 100644
--- a/README.md
+++ b/README.md
@@ -1,3 +1,9 @@
+# Getting started
+
+    virtualenv -p python3 env
+    env/bin/pip install -r requirements.txt
+    PYTHONPATH=. env/bin/python -m ee.tools.read_ltspice_raw demo/reverse-polarity-bak/reverse-polarity_normal.raw
+
 # LTSpice binary format
 
 ## 'real' data
diff --git a/ee/__init__.py b/ee/__init__.py
index 5c26c0a..08a283c 100644
--- a/ee/__init__.py
+++ b/ee/__init__.py
@@ -11,10 +11,17 @@ class LtSpiceRaw(object):
     self.values = values
 
   def get_values(self, variable):
-    v = [v for v in self.variables if v.expression == variable]
-    if len(v) != 1:
-        raise Error('Unknown variable: ' + str(variable))
-    return self.values[v[0].idx]
+    return self.values[variable.idx]
+
+  def get_variable(self, idx=None, expression=None):
+    if idx is not None:
+      return self.variables[idx]
+    if expression is not None:
+      v = [v for v in self.variables if v.expression == expression]
+      if len(v) != 1:
+          raise Exception('Unknown variable: ' + str(variable))
+      return v[0]
+    raise Exception('idx or expression must be given')
 
 class LtSpiceVariable(object):
   def __init__(self, idx, expression, kind):
@@ -28,6 +35,8 @@ class LtSpiceReader(object):
     self.f = f
     self.mode = 'header'
     self.headers = {}
+    self.variables = []
+    self.values = []
 
   def read_header(self, line):
     sep = line.find(':')
@@ -40,7 +49,6 @@ class LtSpiceReader(object):
       self.set_binary_mode()
     elif key == 'Variables':
       self.no_variables = int(self.headers['No. Variables'])
-      self.variables = []
       self.mode = 'variables'
     else:
       self.headers[key] = value
@@ -60,8 +68,6 @@ class LtSpiceReader(object):
 
   def read_binary(self):
     pos = self.f.tell()
-    # Skip an extra '\0'
-    self.f.seek(pos + 1)
     no_points = int(self.headers['No. Points'])
     no_variables = int(self.headers['No. Variables'])
 #    print("variables={}, no_points={}, pos={}".format(no_variables, no_points, pos))
@@ -82,12 +88,14 @@ class LtSpiceReader(object):
       if len(line) == 0:
         break
 
-      if line[0] == 0:
-        line = line[1:]
-      if line[-1] == 10:
+      self.f.read(1) # The data is utf16 encoded so read the extra null byte
+
+      # if this wasn't the last line .readline() includes the newline
+      if line[-1] == '\n' or line[-1] == 0x0a:
         line = line[:-1]
-#    print("len(line)={}, line={}".format(len(line), str(line)))
-      line = str(line, encoding="utf-16")
+
+      print("len(line)={}, line={}".format(len(line), str(line)))
+      line = line.decode(encoding="utf-16")
 
       if self.mode == 'header':
         self.read_header(line)
@@ -100,6 +108,9 @@ class LtSpiceReader(object):
         self.read_binary()
         break
 
+    if len(self.variables) == 0:
+      raise Exception("Something didn't quite work when parsing file, no variables found")
+
     return LtSpiceRaw(self.variables, self.values)
 
 def read_ltspice_raw(filename):
diff --git a/ee/tools/read_ltspice_raw.py b/ee/tools/read_ltspice_raw.py
index 0b3b7e1..279fe73 100644
--- a/ee/tools/read_ltspice_raw.py
+++ b/ee/tools/read_ltspice_raw.py
@@ -12,7 +12,9 @@ raw = ee.read_ltspice_raw(sys.argv[1])
 #  for p in raw.values[i]:
 #    print("  {}".format(p))
 
-xs = raw.get_values('V(load)')
+x = raw.get_variable(idx = 0)
+y = raw.get_variable(expression = 'V(load)')
+xs = raw.get_values(y)
 
 import matplotlib
 matplotlib.use('Agg') 
@@ -22,6 +24,8 @@ from matplotlib.ticker import FuncFormatter, MaxNLocator
 fig = plt.figure()
 ax = fig.add_subplot(111)
 ax.plot(xs)
+ax.set_xlabel(x.expression)
+ax.set_ylabel(y.expression)
 
 with open("ltspice.png", "wb") as f:
   plt.savefig(f, format="png")