From 769d45b7ad7f00667d4053704793d3220c4bb25f Mon Sep 17 00:00:00 2001
From: Ross Patterson <me@rpatterson.net>
Date: Sat, 10 Aug 2019 04:02:41 -0700
Subject: [PATCH] Detect: Parse multiple boards into model lines and model
 numbers

---
 ffmpeg/_detect.py           | 71 +++++++++++++++++++++++++++++++++++++
 ffmpeg/tests/test_ffmpeg.py | 22 ++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/ffmpeg/_detect.py b/ffmpeg/_detect.py
index c88644b..bf22797 100644
--- a/ffmpeg/_detect.py
+++ b/ffmpeg/_detect.py
@@ -39,6 +39,9 @@ parser.add_argument(
     '--ffmpeg', default='ffmpeg',
     help='The path to the ffmpeg execuatble')
 
+# Separators to divide a range of models within a line
+MODEL_RANGE_SEPARATORS = ['-', '>']
+
 # List `hwaccel` options by order of expected performance when available.
 HWACCELS_BY_PERFORMANCE = [
     # NVidia
@@ -106,6 +109,18 @@ def detect_gpus():
 
     if not gpus:
         raise ValueError('No GPUs detected')
+
+    data = _get_data()
+    for gpu in gpus:
+        vendor_data = data.get(gpu.get('vendor', '').lower())
+        if vendor_data:
+            model_lines_data = _parse_models(
+                model_lines=vendor_data['lines'],
+                boards=gpu['board'].lower(), model_data={})
+            gpu['model_line'] = list(model_lines_data.keys())[0]
+            gpu['model_num'] = list(model_lines_data[
+                gpu['model_line']]['models'].keys())[0]
+
     return gpus
 
 
@@ -221,6 +236,62 @@ def _get_data():
     return DATA
 
 
+def _parse_models(
+        model_lines, boards, model_data,
+        model_lines_data=None, model_line=None):
+    """
+    Parse model lines, sets and ranges from a boards string.
+    """
+    if model_lines_data is None:
+        model_lines_data = {}
+
+    boards = boards.strip().lower()
+    model_line_positions = [
+        (boards.index(next_model_line), idx, next_model_line)
+        for idx, next_model_line in enumerate(model_lines)
+        if next_model_line in boards]
+    if model_line_positions:
+        pos, idx, next_model_line = min(model_line_positions)
+        model_group, next_boards = boards.split(next_model_line.lower(), 1)
+    else:
+        model_group = boards
+        next_boards = ''
+
+    model_group = model_group.strip()
+    if model_group:
+        # First item is a model range for the previous model line
+        model_line_data = model_lines_data.setdefault(
+            model_line, dict(models={}, model_ranges=[]))
+
+        models = []
+        for model_split in model_group.split('/'):
+            models.extend(
+                model.strip()
+                for model in model_split.split('+'))
+
+        for model_range in models:
+            for model_range_separator in MODEL_RANGE_SEPARATORS:
+                model_range_parameters = model_range.split(
+                    model_range_separator)
+                if len(model_range_parameters) > 1:
+                    # This is a range of models
+                    model_line_data['model_ranges'].append(
+                        [model_range, model_data])
+                    models.remove(model_range)
+                    break
+            else:
+                model_line_data['models'][model_range] = model_data
+
+    next_boards = next_boards.strip()
+    if next_boards:
+        return _parse_models(
+            model_lines=model_lines, boards=next_boards,
+            model_data=model_data, model_lines_data=model_lines_data,
+            model_line=next_model_line)
+
+    return model_lines_data
+
+
 def main(args=None):
     """
     Dump all ffmpeg build data to json.
diff --git a/ffmpeg/tests/test_ffmpeg.py b/ffmpeg/tests/test_ffmpeg.py
index 6ecf6bb..e27b4e6 100644
--- a/ffmpeg/tests/test_ffmpeg.py
+++ b/ffmpeg/tests/test_ffmpeg.py
@@ -769,6 +769,28 @@ def test__detect():
             assert isinstance(codec_kwargs['output']['codec'], str)
 
 
+def test__detect_parse_models():
+    """
+    Parse model lines, sets and ranges.
+    """
+    model_lines = ffmpeg._detect._parse_models(
+        model_lines=['geforce rtx', 'geforce gtx', 'geforce gt', 'geforce'],
+        boards=(
+            'GeForce GT 630 > 640 GeForce GTX 650 / 660 '
+            'GeForce GT 740+750'),
+        model_data={})
+    assert 'geforce gt' in model_lines
+    assert 'models' in model_lines['geforce gt']
+    assert '740' in model_lines['geforce gt']['models']
+    assert '750' in model_lines['geforce gt']['models']
+    assert 'model_ranges' in model_lines['geforce gt']
+    assert '630 > 640' in model_lines['geforce gt']['model_ranges'][0]
+    assert 'geforce gtx' in model_lines
+    assert 'models' in model_lines['geforce gtx']
+    assert '650' in model_lines['geforce gtx']['models']
+    assert '660' in model_lines['geforce gtx']['models']
+
+
 def get_filter_complex_input(flt, name):
     m = re.search(r'\[([^]]+)\]{}(?=[[;]|$)'.format(name), flt)
     if m: