Line numbers

schema_salad/python_codegen.py

@@ -170,6 +170,8 @@ def begin_class(
             self.out.write("    pass\n\n\n")
             return
 
+        field_names.append("_doc")
+
         required_field_names = [f for f in field_names if f not in optional_fields]
         optional_field_names = [f for f in field_names if f in optional_fields]
 
@@ -274,10 +276,16 @@ def fromDoc(
         self.serializer.write(
             """
     def save(
-        self, top: bool = False, base_url: str = "", relative_uris: bool = True
-    ) -> Dict[str, Any]:
-        r: Dict[str, Any] = {}
-
+        self, top: bool = False, base_url: str = "", relative_uris: bool = True, line_info: Optional[CommentedMap] = None
+    ) -> CommentedMap:
+        r = CommentedMap()
+        if line_info is not None:
+            self._doc = line_info
+        if (type(self._doc) == CommentedMap):
+            r._yaml_set_line_col(self._doc.lc.line,self._doc.lc.col)
+        line_numbers = get_line_numbers(self._doc)
+        max_len = get_max_line_num(self._doc)
+        cols: Dict[int, int] = {}
         if relative_uris:
             for ef in self.extension_fields:
                 r[prefix_url(ef, self.loadingOptions.vocab)] = self.extension_fields[ef]
@@ -301,6 +309,7 @@ def save(
             self.serializer.write(
                 """
         r["class"] = "{class_}"
+        max_len = add_kv(old_doc=self._doc, new_doc=r, line_numbers=line_numbers, key="class", val=r.get("class"), max_len=max_len, cols=cols)
 """.format(
                     class_=classname
                 )
@@ -394,6 +403,7 @@ def type_loader(
             sub_names: List[str] = list(
                 dict.fromkeys([self.type_loader(i).name for i in type_declaration])
             )
+
             return self.declare_type(
                 TypeDef(
                     "union_of_{}".format("_or_".join(sub_names)),
@@ -565,12 +575,15 @@ def declare_field(
 if self.{safename} is not None:
     u = save_relative_uri(self.{safename}, {baseurl}, {scoped_id}, {ref_scope}, relative_uris)
     r["{fieldname}"] = u
+    max_len = add_kv(old_doc = self._doc, new_doc = r, line_numbers = line_numbers, key = "{key_1}", val = r.get("{key_2}"), max_len = max_len, cols = cols)
 """.format(
                         safename=self.safe_name(name),
                         fieldname=shortname(name).strip(),
                         baseurl=baseurl,
                         scoped_id=fieldtype.scoped_id,
                         ref_scope=fieldtype.ref_scope,
+                        key_1=self.safe_name(name),
+                        key_2=self.safe_name(name),
                     ),
                     8,
                 )
@@ -580,9 +593,16 @@ def declare_field(
                 fmt(
                     """
 if self.{safename} is not None:
-    r["{fieldname}"] = save(
-        self.{safename}, top=False, base_url={baseurl}, relative_uris=relative_uris
+    saved_val = save(
+        self.{safename}, top=False, base_url={baseurl}, relative_uris=relative_uris, doc=self._doc.get("{fieldname}")
     )
+
+    if type(saved_val) == list:
+        if len(saved_val) == 1: # If the returned value is a list of size 1, just save the value in the list
+            saved_val = saved_val[0]
+    r["{fieldname}"] = saved_val
+
+    max_len = add_kv(old_doc = self._doc, new_doc = r, line_numbers = line_numbers, key = "{fieldname}", val = r.get("{fieldname}"), max_len = max_len, cols = cols)
 """.format(
                         safename=self.safe_name(name),
                         fieldname=shortname(name),

schema_salad/python_codegen_support.py

@@ -27,7 +27,7 @@
 
 from rdflib import Graph
 from rdflib.plugins.parsers.notation3 import BadSyntax
-from ruamel.yaml.comments import CommentedMap
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 
 from schema_salad.exceptions import SchemaSaladException, ValidationException
 from schema_salad.fetcher import DefaultFetcher, Fetcher, MemoryCachingFetcher
@@ -189,8 +189,8 @@ def fromDoc(
 
     @abstractmethod
     def save(
-        self, top: bool = False, base_url: str = "", relative_uris: bool = True
-    ) -> Dict[str, Any]:
+        self, top: bool = False, base_url: str = "", relative_uris: bool = True, line_info: Optional[CommentedMap] = None
+    ) -> CommentedMap:
         """Convert this object to a JSON/YAML friendly dictionary."""
 
 
@@ -219,27 +219,110 @@ def load_field(val, fieldtype, baseuri, loadingOptions):
     Union[MutableMapping[str, Any], MutableSequence[Any], int, float, bool, str]
 ]
 
+def add_kv(old_doc: CommentedMap, new_doc: CommentedMap, line_numbers: dict[Any,dict[str,int]], key: str, val: Any, max_len: int, cols: dict[int,int])->int:
+    """Add key value pair into Commented Map.
+
+    Function to add key value pair into new CommentedMap given old CommentedMap, line_numbers for each key/val pair in the old CommentedMap,
+    key/val pair to insert, max_line of the old CommentedMap, and max col value taken for each line.
+    """
+    if key in line_numbers:  # If the key to insert is in the original CommentedMap
+        new_doc.lc.add_kv_line_col(key, old_doc.lc.data[key])
+    elif isinstance(val, (int, float, bool, str)):  # If the value is hashable
+        if val in line_numbers:  # If the value is in the original CommentedMap
+            line = line_numbers[val]["line"]
+            if line in cols:
+                col = max(line_numbers[val]["col"], cols[line])
+            else:
+                col = line_numbers[val]["col"]
+            new_doc.lc.add_kv_line_col(key, [line, col, line, col + len(key) + 2])
+            cols[line] = col + len("id") + 2
+    else:  # If neither the key or value is in the original CommentedMap (or value is not hashable)
+        new_doc.lc.add_kv_line_col(key, [max_len, 0, max_len, len(key) + 2])
+        max_len += 1
+    return max_len
+
+
+def get_line_numbers(doc: CommentedMap) -> dict[Any, dict[str, int]]:
+    """Get line numbers for kv pairs in CommentedMap.
+
+    For each key/value pair in a CommentedMap, save the line/col info into a dictionary,
+    only save value info if value is hashable.
+    """
+    line_numbers: Dict[Any, dict[str,int]] = {}
+    if type(doc) == dict:
+        return {}
+    for key, value in doc.lc.data.items():
+        line_numbers[key] = {}
+
+        line_numbers[key]["line"] = doc.lc.data[key][0]
+        line_numbers[key]["col"] = doc.lc.data[key][1]
+        if isinstance(value, (int, float, bool, str)):
+            line_numbers[value] = {}
+            line_numbers[value]["line"] = doc.lc.data[key][2]
+            line_numbers[value]["col"] = doc.lc.data[key][3]
+    return line_numbers
+
+
+def get_max_line_num(doc: CommentedMap) -> int:
+    """Get the max line number for a CommentedMap.
+
+    Iterate through the the key with the highest line number until you reach a non-CommentedMap value or empty CommentedMap.
+    """
+    max_line = 0
+    max_key = ""
+    cur = doc
+    while type(cur) == CommentedMap and len(cur) > 0:
+        for key in cur.lc.data.keys():
+            if cur.lc.data[key][2] >= max_line:
+                max_line = cur.lc.data[key][2]
+            max_key = key
+        cur = cur[max_key]
+    return max_line + 1
+
 
 def save(
     val: Any,
     top: bool = True,
     base_url: str = "",
     relative_uris: bool = True,
+    doc: Optional[CommentedMap] = None,
 ) -> save_type:
+    """Save a val of any type.
+
+    Recursively calls save method from class if val is of type Saveable. Otherwise, saves val to CommentedMap or CommentedSeq
+    """
     if isinstance(val, Saveable):
-        return val.save(top=top, base_url=base_url, relative_uris=relative_uris)
+        return val.save(top=top, base_url=base_url, relative_uris=relative_uris, line_info=doc)
     if isinstance(val, MutableSequence):
-        return [
-            save(v, top=False, base_url=base_url, relative_uris=relative_uris)
-            for v in val
-        ]
+        r = CommentedSeq()
+        for v in val:
+            if doc:
+                if isinstance(v,(int, float, bool, str)):
+                    if v in doc:
+                        r.lc.data.add_kv_line_col(v, doc.lc.data[v])
+            r.append(save(v, top=False, base_url=base_url, relative_uris=relative_uris, doc=doc))
+        return r
+        # return [
+        #     save(v, top=False, base_url=base_url, relative_uris=relative_uris)
+        #     for v in val
+        # ]
     if isinstance(val, MutableMapping):
-        newdict = {}
+        newdict = CommentedMap()
         for key in val:
+            if doc:
+                if isinstance(key, (int, float, bool, str)):
+                    if key in doc:
+                        newdict.lc.add_kv_line_col(key, doc.lc.data[key])
             newdict[key] = save(
-                val[key], top=False, base_url=base_url, relative_uris=relative_uris
+                val[key], top=False, base_url=base_url, relative_uris=relative_uris, doc=doc
             )
         return newdict
+        # newdict = {}
+        # for key in val:
+        #     newdict[key] = save(
+        #         val[key], top=False, base_url=base_url, relative_uris=relative_uris
+        #     )
+        # return newdict
     if val is None or isinstance(val, (int, float, bool, str)):
         return val
     raise Exception("Not Saveable: %s" % type(val))
@@ -710,11 +793,18 @@ def _document_load(
             addl_metadata=addl_metadata,
         )
 
-        doc = {
-            k: v
-            for k, v in doc.items()
-            if k not in ("$namespaces", "$schemas", "$base")
-        }
+        # doc = {
+        #     k: v
+        #     for k, v in doc.items()
+        #     if k not in ("$namespaces", "$schemas", "$base")
+        # }
+
+        if "$namespaces" in doc:
+            doc.pop("$namespaces")
+        if "$schemas" in doc:
+            doc.pop("$schemas")
+        if "$base" in doc:
+            doc.pop("$base")
 
         if "$graph" in doc:
             loadingOptions.idx[baseuri] = (

schema_salad/tests/count-lines6-wf_v1_0.cwl

@@ -0,0 +1,26 @@
+#!/usr/bin/env cwl-runner
+class: Workflow
+cwlVersion: v1.0
+
+requirements:
+  - class: ScatterFeatureRequirement
+  - class: MultipleInputFeatureRequirement
+
+inputs:
+    file1: File[]
+    file2: File[]
+
+outputs:
+    count_output:
+      type: int
+      outputSource: step1/output
+
+steps:
+  step1:
+    run: wc3-tool_v1_0.cwl
+    scatter: file1
+    in:
+      file1:
+        source: [file1, file2]
+        linkMerge: merge_nested
+    out: [output]

schema_salad/tests/count-lines6-wf_v1_1.cwl

@@ -0,0 +1,26 @@
+#!/usr/bin/env cwl-runner
+class: Workflow
+cwlVersion: v1.1
+
+requirements:
+  - class: ScatterFeatureRequirement
+  - class: MultipleInputFeatureRequirement
+
+inputs:
+    file1: File[]
+    file2: File[]
+
+outputs:
+    count_output:
+      type: int
+      outputSource: step1/output
+
+steps:
+  step1:
+    run: wc3-tool_v1_1.cwl
+    scatter: file1
+    in:
+      file1:
+        source: [file1, file2]
+        linkMerge: merge_nested
+    out: [output]

schema_salad/tests/count-lines6-wf_v1_2.cwl

@@ -0,0 +1,27 @@
+
+#!/usr/bin/env cwl-runner
+class: Workflow
+cwlVersion: v1.2
+
+requirements:
+  - class: ScatterFeatureRequirement
+  - class: MultipleInputFeatureRequirement
+
+inputs:
+    file1: File[]
+    file2: File[]
+
+outputs:
+    count_output:
+      type: int
+      outputSource: step1/output
+
+steps:
+  step1:
+    run: wc3-tool_v1_2.cwl
+    scatter: file1
+    in:
+      file1:
+        source: [file1, file2]
+        linkMerge: merge_nested
+    out: [output]

schema_salad/tests/test_line_numbers.py

@@ -0,0 +1,64 @@
+#from parser import load_document_by_uri, save
+from pathlib import Path
+from schema_salad.utils import yaml_no_ts
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
+from typing import Any, Dict, List, Optional, cast
+from schema_salad import codegen
+from schema_salad.avro.schema import Names
+from schema_salad.schema import load_schema
+
+
+def compare_line_numbers(original_doc:CommentedMap, codegen_doc:CommentedMap)->None:
+    assert type(original_doc) == CommentedMap
+    assert type(codegen_doc) == CommentedMap
+
+    assert original_doc.lc.line == codegen_doc.lc.line
+    assert original_doc.lc.col == codegen_doc.lc.col
+
+    for key, lc_info in original_doc.lc.data.items():
+        assert key in codegen_doc.lc.data
+        assert lc_info==codegen_doc.lc.data[key]
+
+    max_line = get_max_line_number(original_doc)
+
+    for key, lc_info in codegen_doc.lc.data.items():
+        if key in original_doc:
+            continue
+        assert lc_info == [max_line, 0, max_line, len(key) + 2]
+        max_line += 1
+
+def get_max_line_number(original_doc:CommentedMap)->int:
+    max_key = ""
+    max_line = 0
+    temp_doc = original_doc
+    while (type(temp_doc) == CommentedMap) and len(temp_doc) > 0:
+        for key, lc_info in temp_doc.lc.data.items():
+            if lc_info[0] >= max_line:
+                max_line = lc_info[0]
+                max_key = key
+        temp_doc = temp_doc[max_key]
+    return max_line + 1
+
+def python_codegen(
+    file_uri: str,
+    target: Path,
+    parser_info: Optional[str] = None,
+    package: Optional[str] = None,
+) -> None:
+    document_loader, avsc_names, schema_metadata, metaschema_loader = load_schema(
+        file_uri
+    )
+    assert isinstance(avsc_names, Names)
+    schema_raw_doc = metaschema_loader.fetch(file_uri)
+    schema_doc, schema_metadata = metaschema_loader.resolve_all(
+        schema_raw_doc, file_uri
+    )
+    codegen.codegen(
+        "python",
+        cast(List[Dict[str, Any]], schema_doc),
+        schema_metadata,
+        document_loader,
+        target=str(target),
+        parser_info=parser_info,
+        package=package
+    )