Merge pull request #985 from SciML/myb/write

Overload `write` on `AbstractSystem`s to write formatted Julia expressions

Author: YingboMa

Project.toml

@@ -15,6 +15,7 @@ Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
@@ -50,6 +51,7 @@ DiffRules = "0.1, 1.0"
 Distributions = "0.23, 0.24"
 DocStringExtensions = "0.7, 0.8"
 IfElse = "0.1"
+JuliaFormatter = "0.13"
 LabelledArrays = "1.3"
 Latexify = "0.11, 0.12, 0.13, 0.14, 0.15"
 LightGraphs = "1.3"

src/ModelingToolkit.jl

@@ -35,6 +35,7 @@ import SymbolicUtils: istree, arguments, operation, similarterm, promote_symtype
 using SymbolicUtils.Code
 import SymbolicUtils.Code: toexpr
 import SymbolicUtils.Rewriters: Chain, Postwalk, Prewalk, Fixpoint
+import JuliaFormatter
 
 using Reexport
 @reexport using Symbolics

src/systems/abstractsystem.jl

@@ -468,6 +468,9 @@ function push_defaults!(stmt, defs, var2name)
     return defs_name
 end
 
+###
+### System I/O
+###
 function toexpr(sys::AbstractSystem)
     sys = flatten(sys)
     expr = Expr(:block)
@@ -503,6 +506,8 @@ function toexpr(sys::AbstractSystem)
     striplines(expr) # keeping the line numbers is never helpful
 end
 
+Base.write(io::IO, sys::AbstractSystem) = write(io, readable_code(toexpr(sys)))
+
 function Base.show(io::IO, ::MIME"text/plain", sys::AbstractSystem)
     eqs = equations(sys)
     if eqs isa AbstractArray

src/utils.jl

@@ -84,3 +84,23 @@ function indepvar2depvar(s::Sym, args...)
     ns = Sym{T}(nameof(s))(args...)
     @set! ns.metadata = s.metadata
 end
+
+function _readable_code(ex)
+    ex isa Expr || return ex
+    if ex.head === :call
+        f, args = ex.args[1], ex.args[2:end]
+        if f isa Function && (nf = nameof(f); Base.isoperator(nf))
+            expr = Expr(:call, nf)
+            for a in args
+                push!(expr.args, _readable_code(a))
+            end
+            return expr
+        end
+    end
+    expr = Expr(ex.head)
+    for a in ex.args
+        push!(expr.args, _readable_code(a))
+    end
+    expr
+end
+readable_code(expr) = JuliaFormatter.format_text(string(Base.remove_linenums!(_readable_code(expr))))

test/serialization.jl

@@ -19,3 +19,9 @@ for prob in [
 
     run(`$(Base.julia_cmd()) -e $(_cmd)`)
 end
+
+include("../examples/rc_model.jl")
+io = IOBuffer()
+write(io, rc_model)
+sys = include_string(@__MODULE__, String(take!(io)))
+@test sys == flatten(rc_model)